--- language: - en tags: - information-retrieval - query-embedding - toolbench - dynamic-scaling - api-retrieval - tool-discovery license: mit --- # Librarian of Tools A sophisticated query embedding model designed to act as an intelligent "librarian" for discovering and retrieving the most relevant tools and APIs based on natural language queries. ## 🎯 Purpose This model serves as a **Librarian of Tools** - an AI system that understands user intentions and finds the most appropriate tools, APIs, or functions to accomplish their tasks. It's particularly effective for: - **API Discovery**: Finding relevant rapid APIs from large collections - **Iterative Search**: Progressive tool discovery with residual-based refinement ## 🏗️ Model Architecture - **Base Model**: ToolBench/ToolBench_IR_bert_based_uncased - **Architecture**: Query embedding with scale prediction - **Special Features**: - **Query-Specific Adaptation**: Tailors embeddings to individual query characteristics - **Balanced Magnitude Handling**: Maintains appropriate scaling for retrieval tasks - **Residual-Based Iteration**: Supports iterative search for comprehensive tool discovery ## 🎓 Training Strategy - **Training Approach**: Dynamic direction-focused with AdamW optimizer - **Loss Function**: Combined MSE, direction loss, and magnitude loss - **Scale Prediction**: Uses softplus + 1 activation for 1 + scale factors - **Dataset**: Trained on query-(sum vector of relevant API embeddings) pairs from ToolBench(https://github.com/OpenBMB/ToolBench) ## 🚀 Usage ### Basic Usage ```python from transformers import AutoTokenizer, AutoModel import torch # Load the Librarian model model = AutoModel.from_pretrained("jhleepidl/librarian") tokenizer = AutoTokenizer.from_pretrained("jhleepidl/librarian") # Ask the librarian to find tools query = "How to send an email using Python?" inputs = tokenizer(query, return_tensors="pt", truncation=True, max_length=256) # Get the embedding for tool retrieval with torch.no_grad(): embedding = model(**inputs) # The output is optimized for finding relevant tools/APIs ``` ### Advanced Usage with Similarity Search ```python import numpy as np from sklearn.metrics.pairwise import cosine_similarity # Get embeddings for multiple queries queries = [ "Send email with attachment", "Download file from URL", "Parse JSON data", "Make HTTP POST request" ] embeddings = [] for query in queries: inputs = tokenizer(query, return_tensors="pt", truncation=True, max_length=256) with torch.no_grad(): embedding = model(**inputs).cpu().numpy() embeddings.append(embedding.flatten()) # Find similar queries similarity_matrix = cosine_similarity(embeddings) print("Query similarity matrix:") print(similarity_matrix) ``` ### Iterative Search for Comprehensive Tool Discovery The model supports iterative search, which progressively discovers tools by removing found APIs from the query representation and continuing the search: ```python import torch import torch.nn.functional as F import numpy as np class LibrarianSearch: def __init__(self, model, tokenizer, vector_db_index, documents, threshold=0.5): self.model = model self.tokenizer = tokenizer self.index = vector_db_index self.documents = documents self.threshold = threshold def get_query_embedding(self, query, normalize=False): """Get query embedding using the Librarian model""" inputs = self.tokenizer( query, return_tensors="pt", truncation=True, max_length=256 ) with torch.no_grad(): embedding = self.model(**inputs) embedding = embedding.cpu().numpy() if normalize: embedding = embedding / np.linalg.norm(embedding) return embedding.flatten() def iterative_greedy_search(self, query, remove_duplicates=True): """ Perform iterative greedy search to find multiple relevant tools Args: query: The search query remove_duplicates: Whether to avoid returning duplicate APIs Returns: List of found APIs with scores """ # Get initial query embedding (unnormalized for residual calculation) query_embedding_unnorm = self.get_query_embedding(query, normalize=False) query_embedding_norm = self.get_query_embedding(query, normalize=True) found_apis = [] current_query_unnorm = query_embedding_unnorm.copy() current_query_norm = query_embedding_norm.copy() found_api_keys = set() if remove_duplicates else None while True: # Search for the best matching API scores, indices = self.index.search( current_query_norm.reshape(1, -1), 1 ) if indices[0][0] == -1 or scores[0][0] < self.threshold: break # Get the found API idx = indices[0][0] doc = self.documents[idx] api_key = f"{doc['metadata']['tool_name']}_{doc['metadata']['api_name']}" # Check for duplicates if remove_duplicates and api_key in found_api_keys: # Calculate residual and continue doc_embedding = self.get_embedding_by_index(idx) residual = current_query_unnorm - doc_embedding residual_norm = np.linalg.norm(residual) if residual_norm > np.linalg.norm(current_query_unnorm): break current_query_unnorm = residual if residual_norm > 0: current_query_norm = residual / residual_norm else: break continue # Add the found API found_apis.append({ 'tool_name': doc['metadata']['tool_name'], 'api_name': doc['metadata']['api_name'], 'score': float(scores[0][0]) }) if remove_duplicates: found_api_keys.add(api_key) # Calculate residual and update query doc_embedding = self.get_embedding_by_index(idx) residual = current_query_unnorm - doc_embedding residual_norm = np.linalg.norm(residual) if residual_norm > np.linalg.norm(current_query_unnorm): break current_query_unnorm = residual if residual_norm > 0: current_query_norm = residual / residual_norm else: break return found_apis def get_embedding_by_index(self, idx): """Get embedding for a specific document index""" # Implementation depends on your vector database setup return self.index.reconstruct(int(idx)) # Usage example librarian_search = LibrarianSearch(model, tokenizer, vector_db_index, documents) # Find multiple tools for a complex query query = "I need to send emails, process images, and analyze data" found_tools = librarian_search.iterative_greedy_search(query) print("Found tools:") for tool in found_tools: print(f"- {tool['tool_name']}.{tool['api_name']} (score: {tool['score']:.3f})") ``` ### Beam Search for Optimal Tool Combinations For even more sophisticated tool discovery, the model supports beam search: ```python def beam_search_iterative(self, query, beam_size=5): """ Perform beam search to find optimal combinations of tools Args: query: The search query beam_size: Number of beams to maintain Returns: Best combination of APIs found """ # Implementation similar to iterative_greedy_search but maintains # multiple candidate paths (beams) and selects the best combination # This is useful for complex queries requiring multiple tools # ... beam search implementation ... pass ``` ## 📚 Citation If you use this model in your research or applications, please cite: ```bibtex @misc{librarian_of_tools, title={Librarian of Tools}, author={jhleepidl}, year={2025}, url={https://github.com/jhleepidl/librarian} } ``` ## 📄 License This model is released under the MIT License, making it suitable for both research and commercial applications. --- **The Librarian of Tools** - Your intelligent assistant for discovering the right tools for any task! 🛠️📚