Views
No views yet
1from transformers import AutoTokenizer, AutoModel
2import torch
3
4# Load the Librarian model
5model = AutoModel.from_pretrained("jhleepidl/librarian")
6tokenizer = AutoTokenizer.from_pretrained("jhleepidl/librarian")
7
8# Ask the librarian to find tools
9query = "How to send an email using Python?"
10inputs = tokenizer(query, return_tensors="pt", truncation=True, max_length=256)
11
12# Get the embedding for tool retrieval
13with torch.no_grad():
14 embedding = model(**inputs)
15 # The output is optimized for finding relevant tools/APIs1import numpy as np
2from sklearn.metrics.pairwise import cosine_similarity
3
4# Get embeddings for multiple queries
5queries = [
6 "Send email with attachment",
7 "Download file from URL",
8 "Parse JSON data",
9 "Make HTTP POST request"
10]
11
12embeddings = []
13for query in queries:
14 inputs = tokenizer(query, return_tensors="pt", truncation=True, max_length=256)
15 with torch.no_grad():
16 embedding = model(**inputs).cpu().numpy()
17 embeddings.append(embedding.flatten())
18
19# Find similar queries
20similarity_matrix = cosine_similarity(embeddings)
21print("Query similarity matrix:")
22print(similarity_matrix)1import torch
2import torch.nn.functional as F
3import numpy as np
4
5class LibrarianSearch:
6 def __init__(self, model, tokenizer, vector_db_index, documents, threshold=0.5):
7 self.model = model
8 self.tokenizer = tokenizer
9 self.index = vector_db_index
10 self.documents = documents
11 self.threshold = threshold
12
13 def get_query_embedding(self, query, normalize=False):
14 """Get query embedding using the Librarian model"""
15 inputs = self.tokenizer(
16 query,
17 return_tensors="pt",
18 truncation=True,
19 max_length=256
20 )
21
22 with torch.no_grad():
23 embedding = self.model(**inputs)
24 embedding = embedding.cpu().numpy()
25
26 if normalize:
27 embedding = embedding / np.linalg.norm(embedding)
28 return embedding.flatten()
29
30 def iterative_greedy_search(self, query, remove_duplicates=True):
31 """
32 Perform iterative greedy search to find multiple relevant tools
33
34 Args:
35 query: The search query
36 remove_duplicates: Whether to avoid returning duplicate APIs
37
38 Returns:
39 List of found APIs with scores
40 """
41 # Get initial query embedding (unnormalized for residual calculation)
42 query_embedding_unnorm = self.get_query_embedding(query, normalize=False)
43 query_embedding_norm = self.get_query_embedding(query, normalize=True)
44
45 found_apis = []
46 current_query_unnorm = query_embedding_unnorm.copy()
47 current_query_norm = query_embedding_norm.copy()
48 found_api_keys = set() if remove_duplicates else None
49
50 while True:
51 # Search for the best matching API
52 scores, indices = self.index.search(
53 current_query_norm.reshape(1, -1), 1
54 )
55
56 if indices[0][0] == -1 or scores[0][0] < self.threshold:
57 break
58
59 # Get the found API
60 idx = indices[0][0]
61 doc = self.documents[idx]
62 api_key = f"{doc['metadata']['tool_name']}_{doc['metadata']['api_name']}"
63
64 # Check for duplicates
65 if remove_duplicates and api_key in found_api_keys:
66 # Calculate residual and continue
67 doc_embedding = self.get_embedding_by_index(idx)
68 residual = current_query_unnorm - doc_embedding
69 residual_norm = np.linalg.norm(residual)
70
71 if residual_norm > np.linalg.norm(current_query_unnorm):
72 break
73
74 current_query_unnorm = residual
75 if residual_norm > 0:
76 current_query_norm = residual / residual_norm
77 else:
78 break
79 continue
80
81 # Add the found API
82 found_apis.append({
83 'tool_name': doc['metadata']['tool_name'],
84 'api_name': doc['metadata']['api_name'],
85 'score': float(scores[0][0])
86 })
87
88 if remove_duplicates:
89 found_api_keys.add(api_key)
90
91 # Calculate residual and update query
92 doc_embedding = self.get_embedding_by_index(idx)
93 residual = current_query_unnorm - doc_embedding
94 residual_norm = np.linalg.norm(residual)
95
96 if residual_norm > np.linalg.norm(current_query_unnorm):
97 break
98
99 current_query_unnorm = residual
100 if residual_norm > 0:
101 current_query_norm = residual / residual_norm
102 else:
103 break
104
105 return found_apis
106
107 def get_embedding_by_index(self, idx):
108 """Get embedding for a specific document index"""
109 # Implementation depends on your vector database setup
110 return self.index.reconstruct(int(idx))
111
112# Usage example
113librarian_search = LibrarianSearch(model, tokenizer, vector_db_index, documents)
114
115# Find multiple tools for a complex query
116query = "I need to send emails, process images, and analyze data"
117found_tools = librarian_search.iterative_greedy_search(query)
118
119print("Found tools:")
120for tool in found_tools:
121 print(f"- {tool['tool_name']}.{tool['api_name']} (score: {tool['score']:.3f})")1def beam_search_iterative(self, query, beam_size=5):
2 """
3 Perform beam search to find optimal combinations of tools
4
5 Args:
6 query: The search query
7 beam_size: Number of beams to maintain
8
9 Returns:
10 Best combination of APIs found
11 """
12 # Implementation similar to iterative_greedy_search but maintains
13 # multiple candidate paths (beams) and selects the best combination
14 # This is useful for complex queries requiring multiple tools
15
16 # ... beam search implementation ...
17 pass1@misc{librarian_of_tools,
2 title={Librarian of Tools},
3 author={jhleepidl},
4 year={2025},
5 url={https://github.com/jhleepidl/librarian}
6}