Batch generate embeddings in add_nodes / add_edges#2201
Conversation
fabubaker
changed the title
add_nodesadd_nodes / add_edges
| let vectors = self.cache.get_embeddings(docs).await?; | ||
|
|
||
| for (id, vector) in ids.iter().zip(vectors) { | ||
| self.edge_db.insert_vector(*id, &vector)?; |
There was a problem hiding this comment.
Every time you call insert_vector this is going to update the index. We should probably change that function to accept a vector of embeddings instead, so we can write everything to arroy and only in the end update the index just once. The same applies to nodes
fabubaker
force-pushed
the
feat/batch-create-embeddings
branch
from
7e268b1 to
0dcf33e
Compare
|
|
ricopinazo
left a comment
ricopinazo
left a comment
There was a problem hiding this comment.
LGTM, just left a small comment that you might want to have a look at, but it's not the end of the world
| } | ||
|
|
||
| // Generate embeddings | ||
| let edges: Vec<_> = edges.into_iter().collect::<Result<Vec<_>, _>>()?; |
There was a problem hiding this comment.
I'm not sure if might be posible to collect here just once
There was a problem hiding this comment.
The first collect uses ? to return the first error it encounters in edges, so both collect are needed.
That said I did collapse both into one chain and got rid of the intermediate variable.
| type GraphAlgorithmPlugin { | ||
| shortest_path(source: String!, targets: [String!]!, direction: String): [ShortestPathOutput!]! | ||
| pagerank(iterCount: Int!, threads: Int, tol: Float): [PagerankOutput!]! | ||
| shortest_path(source: String!, targets: [String!]!, direction: String): [ShortestPathOutput!]! |
There was a problem hiding this comment.
Just letting a comment here to remember we should make sure whatever generates de GraphQL schema file should do it in a predictable way to avoid these kinds of conflicts
…htory into feat/batch-create-embeddings
3 participants
What changes were proposed in this pull request?
Currently,
GqlMutableGraph.add_nodes/GqlMutableGraph.add_edgesrepeatedly calls the embedding function for each node/edge. This PR modifies it to use the function just once for an entire batch of nodes/edges.Why are the changes needed?
Optimizes how embeddings are created, for eg, by using a batch call to OpenAI.
Does this PR introduce any user-facing change? If yes is this documented?
No.
How was this patch tested?
Unit tests were added to
mutable_graph.rs.Are there any further changes required?
No.