| 8 | device = "cuda" # for GPU usage or "cpu" for CPU usage |
| 9 | |
| 10 | class CustomModel(Model): |
| 11 | def __init__(self, name: str, org_name: str, repo_name: str, max_sequence_length: int): |
| 12 | super().__init__(name) |
| 13 | self.name = name |
| 14 | self.org_name = org_name |
| 15 | self.repo_name = repo_name |
| 16 | self.max_sequence_length = max_sequence_length |
| 17 | self.device = torch.device(device if torch.cuda.is_available() else "cpu") |
| 18 | print("Using device:", self.device) |
| 19 | self.ready = False |
| 20 | |
| 21 | def load(self): |
| 22 | self.tokenizer = AutoTokenizer.from_pretrained(f"{self.org_name}/{self.repo_name}", trust_remote_code=True) |
| 23 | self.model = SentenceTransformer(f"{self.org_name}/{self.repo_name}", trust_remote_code=True) |
| 24 | self.model.max_seq_length = self.max_sequence_length |
| 25 | self.ready = True |
| 26 | |
| 27 | def predict(self, payload: Dict, headers: Dict) -> Dict: |
| 28 | inputs = payload["instances"] |
| 29 | return self.process_multiple_files(inputs) |
| 30 | |
| 31 | def process_multiple_files(self, file_inputs: List[Dict]) -> Dict: |
| 32 | all_chunks = [] |
| 33 | file_chunk_map = {} |
| 34 | |
| 35 | # Accumulate chunks from all files |
| 36 | for file_input in file_inputs: |
| 37 | file_path = file_input["file_path"] |
| 38 | source_code = file_input["code"] |
| 39 | hash = file_input["file_hash"] |
| 40 | chunks = self.chunk_code(source_code, hash, 500) |
| 41 | all_chunks.extend(chunks) |
| 42 | file_chunk_map[file_path] = (len(all_chunks) - len(chunks), len(all_chunks)) |
| 43 | |
| 44 | # Encode all chunks at once |
| 45 | codes = [chunk["code"] for chunk in all_chunks] |
| 46 | code_embs = self.model.encode(codes, convert_to_tensor=True) |
| 47 | |
| 48 | # Distribute embeddings back to respective files |
| 49 | results = {} |
| 50 | for file_path, (start, end) in file_chunk_map.items(): |
| 51 | file_chunks = all_chunks[start:end] |
| 52 | for chunk, code_emb in zip(file_chunks, code_embs[start:end]): |
| 53 | chunk["embedding"] = code_emb.tolist() |
| 54 | results[file_path] = {"embeddings": file_chunks} |
| 55 | |
| 56 | return {"results": results} |
| 57 | |
| 58 | def chunk_code(self, code, hash, max_token_length): |
| 59 | # Encode the entire code at once, ignoring special tokens |
| 60 | tokens_data = self.tokenizer.encode_plus(code, add_special_tokens=False, return_offsets_mapping=True) |
| 61 | tokens = tokens_data['input_ids'] |
| 62 | offsets = tokens_data['offset_mapping'] |
| 63 | |
| 64 | chunks = [] |
| 65 | current_chunk_start_index = 0 |
| 66 | total_tokens = len(tokens) |
| 67 | current_token_index = 0 |