Find Related: Save embeddings in cache

This optimizes initialization on the virtual machine

routers/tool_find_related.py

@@ -1,5 +1,7 @@
 # find_related.py
 
+import os
+import pickle
 import re
 import torch
 import threading
@@ -11,7 +13,6 @@ try:
     from .utils_gitea import gitea_fetch_issues, gitea_json_issue_get
     from config import settings
 except:
-    import os
     import sys
     from utils_gitea import gitea_fetch_issues, gitea_json_issue_get
     sys.path.append(os.path.abspath(
@@ -58,6 +59,7 @@ class EmbeddingContext:
     TOKEN_LEN_MAX_FOR_EMBEDDING = 512
     TOKEN_LEN_MAX_BALCKLIST = 2 * TOKEN_LEN_MAX_FOR_EMBEDDING
     issue_attr_filter = {'number', 'title', 'body', 'state', 'updated_at'}
+    cache_path = "routers/tool_find_related_cache.pkl"
 
     # Set when creating the object
     lock = None
@@ -68,8 +70,8 @@ class EmbeddingContext:
 
     # Updates constantly
     data = {}
-    black_list = {'blender': {'blender': {109399, 113157, 114706},
-                              'blender-addons': set()}}
+    black_list = {'blender': {109399, 113157, 114706},
+                  'blender-addons': set()}
 
     def __init__(self):
         self.lock = threading.Lock()
@@ -165,15 +167,19 @@ class EmbeddingContext:
 
         return texts_to_embed
 
-    def embeddings_generate(self, owner, repo):
-        if not owner in self.black_list:
-            self.black_list[owner] = {repo: {}}
-        elif not repo in self.black_list[owner]:
-            self.black_list[owner][repo] = {}
+    def embeddings_generate(self, repo):
+        if os.path.exists(self.cache_path):
+            with open(self.cache_path, 'rb') as file:
+                self.data = pickle.load(file)
+                if repo in self.data:
+                    return
 
-        black_list = self.black_list[owner][repo]
+        if not repo in self.black_list:
+            self.black_list[repo] = {}
 
-        issues = gitea_fetch_issues(owner, repo, state='open', since=None,
+        black_list = self.black_list[repo]
+
+        issues = gitea_fetch_issues('blender', repo, state='open', since=None,
                                     issue_attr_filter=self.issue_attr_filter, exclude=black_list)
 
         issues = sorted(issues, key=lambda issue: int(issue['number']))
@@ -190,26 +196,21 @@ class EmbeddingContext:
             'embeddings': embeddings,
         }
 
-        if not owner in self.data:
-            self.data[owner] = {repo: {}}
-        elif not repo in self.data[owner]:
-            self.data[owner][repo] = {}
-
-        self.data[owner][repo] = data
+        self.data[repo] = data
 
-    def embeddings_updated_get(self, owner, repo):
+    def embeddings_updated_get(self, repo):
         with self.lock:
             try:
-                data = self.data[owner][repo]
+                data = self.data[repo]
             except:
-                self.embeddings_generate(owner, repo)
-                data = self.data[owner][repo]
+                self.embeddings_generate(repo)
+                data = self.data[repo]
 
-            black_list = self.black_list[owner][repo]
+            black_list = self.black_list[repo]
             date_old = data['updated_at']
 
             issues = gitea_fetch_issues(
-                owner, repo, since=date_old, issue_attr_filter=self.issue_attr_filter, exclude=black_list)
+                'blender', repo, since=date_old, issue_attr_filter=self.issue_attr_filter, exclude=black_list)
 
             # Get the most recent date
             date_new = _find_latest_date(issues, date_old)
@@ -361,13 +362,13 @@ def text_search(owner, repo, text_to_embed, limit=None):
 
 
 def find_relatedness(gitea_issue, limit=20):
-    owner = gitea_issue['repository']['owner']
+    assert gitea_issue['repository']['owner'] == 'blender'
     repo = gitea_issue['repository']['name']
     title = gitea_issue['title']
     body = gitea_issue['body']
     number = int(gitea_issue['number'])
 
-    data = EMBEDDING_CTX.embeddings_updated_get(owner, repo)
+    data = EMBEDDING_CTX.embeddings_updated_get(repo)
     new_embedding = None
 
     # Check if the embedding already exist.
@@ -394,38 +395,40 @@ def find_relatedness(gitea_issue, limit=20):
     return '\n'.join(duplicates)
 
 
-@router.get("/find_related")
-def find_related(owner: str = 'blender', repo: str = 'blender', number: int = 1, limit: int = 50):
-    issue = gitea_json_issue_get(owner, repo, number)
+@router.get("/find_related/{repo}/{number}")
+def find_related(repo: str = 'blender', number: int = 104399, limit: int = 50):
+    issue = gitea_json_issue_get('blender', repo, number)
     related = find_relatedness(issue, limit=limit)
     return related
 
 
 if __name__ == "__main__":
-    import os
-    import pickle
-    repo = 'blender-addons'
-    cache_dir = f"routers/cache/{repo}"
-    file_path = os.path.join(cache_dir, "data.pkl")
-
-    if not os.path.exists(cache_dir):
-        os.makedirs(cache_dir, exist_ok=True)
-        with open(file_path, "wb") as file:
-            EMBEDDING_CTX.embeddings_generate('blender', repo)
-            pickle.dump(
-                EMBEDDING_CTX.data['blender'][repo], file, protocol=pickle.HIGHEST_PROTOCOL)
+    update_cache = True
+    if update_cache:
+        EMBEDDING_CTX.embeddings_updated_get('blender')
+        EMBEDDING_CTX.embeddings_updated_get('blender-addons')
+        cache_path = EMBEDDING_CTX.cache_path
+        with open(cache_path, "wb") as file:
+            # Converting the embeddings to be CPU compatible, as the virtual machine in use currently only supports the CPU.
+            for val in EMBEDDING_CTX.data.values():
+                val['embeddings'] = val['embeddings'].to(torch.device('cpu'))
+
+            pickle.dump(EMBEDDING_CTX.data, file,
+                        protocol=pickle.HIGHEST_PROTOCOL)
     else:
-        with open(file_path, 'rb') as file:
-            EMBEDDING_CTX.data['blender'] = {repo: pickle.load(file)}
+        # Converting the embeddings to be GPU.
+        for val in EMBEDDING_CTX.data.values():
+            val['embeddings'] = val['embeddings'].to(torch.device('cuda'))
 
-    # 'blender/blender/111434' must print #96153, #83604 and #79762
-    issue = gitea_json_issue_get('blender', repo, 105027)
-    print(issue['title'])
+        # 'blender/blender/111434' must print #96153, #83604 and #79762
+        issue1 = gitea_json_issue_get('blender', 'blender', 111434)
+        issue2 = gitea_json_issue_get('blender', 'blender-addons', 104399)
 
-    related = find_relatedness(issue, limit=50)
+        related1 = find_relatedness(issue1, limit=20)
+        related2 = find_relatedness(issue2, limit=20)
 
-    if related == '':
-        print("No potential duplicates found.")
-    else:
         print("These are the 20 most related issues:")
-        print(related)
+        print(related1)
+        print()
+        print("These are the 20 most related issues:")
+        print(related2)

routers/tool_find_related_cache.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4278a800fae11df440b487415a95baa437f1b0651b7b007e44aa951d795049a
+size 21250520