| import os |
| import json |
| import torch |
| from sentence_transformers import SentenceTransformer |
| from .utils import get_md5 |
|
|
|
|
| class ToolRAGModel: |
| def __init__(self, rag_model_name): |
| self.rag_model_name = rag_model_name |
| self.rag_model = None |
| self.tool_desc_embedding = None |
| self.tool_name = None |
| self.tool_embedding_path = None |
| self.load_rag_model() |
|
|
| def load_rag_model(self): |
| self.rag_model = SentenceTransformer(self.rag_model_name) |
| self.rag_model.max_seq_length = 4096 |
| self.rag_model.tokenizer.padding_side = "right" |
|
|
| def load_tool_desc_embedding(self, toolbox): |
| self.tool_name, _ = toolbox.refresh_tool_name_desc(enable_full_desc=True) |
| all_tools_str = [json.dumps(each) for each in toolbox.prepare_tool_prompts(toolbox.all_tools)] |
| md5_value = get_md5(str(all_tools_str)) |
| print("Computed MD5 for tool embedding:", md5_value) |
|
|
| self.tool_embedding_path = os.path.join( |
| os.path.dirname(__file__), |
| self.rag_model_name.split("/")[-1] + f"_tool_embedding_{md5_value}.pt" |
| ) |
|
|
| if os.path.exists(self.tool_embedding_path): |
| try: |
| self.tool_desc_embedding = torch.load(self.tool_embedding_path, map_location="cpu") |
| assert len(self.tool_desc_embedding) == len(toolbox.all_tools), \ |
| "Tool count mismatch with loaded embeddings." |
| print("\033[92mLoaded cached tool_desc_embedding.\033[0m") |
| return |
| except Exception as e: |
| print(f"⚠️ Failed loading cached embeddings: {e}") |
| self.tool_desc_embedding = None |
|
|
| print("\033[93mGenerating new tool_desc_embedding...\033[0m") |
| self.tool_desc_embedding = self.rag_model.encode( |
| all_tools_str, prompt="", normalize_embeddings=True |
| ) |
|
|
| torch.save(self.tool_desc_embedding, self.tool_embedding_path) |
| print(f"\033[92mSaved new tool_desc_embedding to {self.tool_embedding_path}\033[0m") |
|
|
| def rag_infer(self, query, top_k=5): |
| torch.cuda.empty_cache() |
| queries = [query] |
| query_embeddings = self.rag_model.encode( |
| queries, prompt="", normalize_embeddings=True |
| ) |
| if self.tool_desc_embedding is None: |
| raise RuntimeError("❌ tool_desc_embedding is not initialized. Did you forget to call load_tool_desc_embedding()?") |
|
|
| scores = self.rag_model.similarity( |
| query_embeddings, self.tool_desc_embedding |
| ) |
| top_k = min(top_k, len(self.tool_name)) |
| top_k_indices = torch.topk(scores, top_k).indices.tolist()[0] |
| top_k_tool_names = [self.tool_name[i] for i in top_k_indices] |
| return top_k_tool_names |
|
|
