refactor update_similar_snippets function
@@ -1288,6 +1288,8 @@ class Client(Interface):
ids, snippets, and embeddings
"""
disc = self.get_discoveries(repo_url)
+ # If called by UI classes, disc is a tuple of 2 elements, and the
+ # actual discoveries are at the second element
discoveries = disc[1] if len(disc) == 2 else disc
discoveries_ids = [d['id'] for d in discoveries]
snippets = [d['snippet'] for d in discoveries]
@@ -1300,6 +1302,7 @@ class Client(Interface):
state,
repo_url,
file_name=None,
+ compute_missing_embeddings=False,
threshold=0.96):
""" Find snippets that are similar to the target
snippet and update their state.
@@ -1314,6 +1317,9 @@ class Client(Interface):
The url of the repository
file_name: str
Restrict to a given file the search for similar snippets
+ compute_missing_embeddings: bool
+ Compute (or not) embeddings when they are missing from the db
+ (default `False`)
threshold: float
Update snippets with similarity score above threshold.
Values lesser than 0.94 do not generally imply any relevant
@@ -1329,25 +1335,42 @@ class Client(Interface):
# Discoveries are the second element of the output of
# get_discoveries in the UI clients, but are the entire
# output in regular clients
+ # (not the elegant way to do it, but hack for double inheritance)
discoveries = disc[1] if len(disc) == 2 else disc
- # Compute target snippet embedding
- target_embedding = self.get_embedding(snippet=target_snippet)
- # TODO: if this snippet has no embedding, compute it (allow the user to
- # chose whether to compute missing embeddings or not)
- # target_snippet_embedding = compute_snippet_embedding(target_snippet,
- # model)
+ # Keep only the discoveries with a state different from the one
+ # passed as argument (same state discoveries will not be updated)
+ discoveries = filter(lambda d: d['state'] != state, discoveries)
- # If specified target snippet is not found in the embeddings table,
- # no update is performed
+ # Get the embedding of the target snippet
+ target_embedding = self.get_embedding(snippet=target_snippet)
+ # Get all embeddings for this repo
+ all_embeddings = self.get_embeddings(repo_url=repo_url)
+ # Check if need to recompute embeddings
+ if not all_embeddings and compute_missing_embeddings:
+ logger.info(f'Compute embeddings for repo {repo_url}')
+ self.add_embeddings(repo_url)
+ all_embeddings = self.get_embeddings(repo_url=repo_url)
if not target_embedding:
+ # It may have just been computed
+ target_embedding = self.get_embedding(snippet=target_snippet)
+
+ # If the target snippet is not found in the embeddings table, or if
+ # the other embeddings are missing, no update is performed
+ if not target_embedding or not all_embeddings:
+ logger.debug('No embeddings found')
return 0
- # We have a target embedding for the target snippet
n_updated_snippets = 0
for d in discoveries:
+ embedding = all_embeddings.get(d['id'])
+ if not embedding and compute_missing_embeddings:
+ # Recompute it
+ logger.debug(f'Compute embedding for discovery {d["id"]}')
+ self.add_embedding(discovery_id=d['id'], repo_url=repo_url)
embedding = self.get_embedding(discovery_id=d['id'])
- if d['state'] != state and embedding:
- # Compute similarity of target embedding and embedding
+ if not embedding:
+ continue
+ # Compute similarity of target_embedding and embedding
similarity = compute_similarity(target_embedding,
embedding)
if similarity > threshold:
