Find Related: Support to closed issues

routers/__init__.py

@@ -0,0 +1 @@
+# __init__.py

routers/tool_find_related.py

@@ -5,7 +5,9 @@ import pickle
 import re
 import torch
 import threading
+
 from datetime import datetime, timedelta
+from enum import Enum
 from sentence_transformers import SentenceTransformer, util
 from fastapi import APIRouter
 
@@ -53,6 +55,7 @@ class EmbeddingContext:
     # These don't change
     TOKEN_LEN_MAX_FOR_EMBEDDING = 512
     TOKEN_LEN_MAX_BALCKLIST = 2 * TOKEN_LEN_MAX_FOR_EMBEDDING
+    ARRAY_CHUNK_SIZE = 4096
     issue_attr_filter = {'number', 'title', 'body', 'state', 'updated_at'}
     cache_path = "routers/tool_find_related_cache.pkl"
 
@@ -62,6 +65,9 @@ class EmbeddingContext:
     openai_client = None
     model_name = ''
     config_type = ''
+    embedding_shape = None
+    embedding_dtype = None
+    embedding_device = None
 
     # Updates constantly
     data = {}
@@ -102,6 +108,11 @@ class EmbeddingContext:
         self.model_name = model_name
         self.config_type = config_type
 
+        tmp = self.encode(['tmp'])
+        self.embedding_shape = tmp.shape[1:]
+        self.embedding_dtype = tmp.dtype
+        self.embedding_device = tmp.device
+
     def encode(self, texts_to_embed):
         pass
 
@@ -171,6 +182,40 @@ class EmbeddingContext:
 
         return texts_to_embed
 
+    def data_ensure_size(self, repo, size_new):
+        updated_at_old = None
+        arrays_size_old = 0
+        titles_old = []
+        try:
+            arrays_size_old = self.data[repo]['arrays_size']
+            if size_new <= arrays_size_old:
+                return
+        except:
+            pass
+
+        arrays_size_new = self.ARRAY_CHUNK_SIZE * \
+            (int(size_new / self.ARRAY_CHUNK_SIZE) + 1)
+
+        data_new = {
+            'updated_at': updated_at_old,
+            'arrays_size': arrays_size_new,
+            'titles': titles_old + [None] * (arrays_size_new - arrays_size_old),
+            'embeddings': torch.empty((arrays_size_new, *self.embedding_shape),
+                                      dtype=self.embedding_dtype,
+                                      device=self.embedding_device),
+            'opened':  torch.zeros(arrays_size_new, dtype=torch.bool),
+            'closed':  torch.zeros(arrays_size_new, dtype=torch.bool),
+        }
+
+        try:
+            data_new['embeddings'][:arrays_size_old] = self.data[repo]['embeddings']
+            data_new['opened'][:arrays_size_old] = self.data[repo]['opened']
+            data_new['closed'][:arrays_size_old] = self.data[repo]['closed']
+        except:
+            pass
+
+        self.data[repo] = data_new
+
     def embeddings_generate(self, repo):
         if os.path.exists(self.cache_path):
             with open(self.cache_path, 'rb') as file:
@@ -183,24 +228,31 @@ class EmbeddingContext:
 
         black_list = self.black_list[repo]
 
-        issues = gitea_fetch_issues('blender', repo, state='open', since=None,
+        issues = gitea_fetch_issues('blender', repo, state='all', since=None,
                                     issue_attr_filter=self.issue_attr_filter, exclude=black_list)
 
-        issues = sorted(issues, key=lambda issue: int(issue['number']))
+        # issues = sorted(issues, key=lambda issue: int(issue['number']))
 
         print("Embedding Issues...")
         texts_to_embed = self.create_strings_to_embbed(issues, black_list)
         embeddings = self.encode(texts_to_embed)
 
-        data = {
-            # Get the most recent date
-            'updated_at': _find_latest_date(issues),
-            'numbers': [int(issue['number']) for issue in issues],
-            'titles': [issue['title'] for issue in issues],
-            'embeddings': embeddings,
-        }
+        self.data_ensure_size(repo, int(issues[0]['number']))
+        self.data[repo]['updated_at'] = _find_latest_date(issues)
+
+        titles = self.data[repo]['titles']
+        embeddings_new = self.data[repo]['embeddings']
+        opened = self.data[repo]['opened']
+        closed = self.data[repo]['closed']
 
-        self.data[repo] = data
+        for i, issue in enumerate(issues):
+            number = int(issue['number'])
+            titles[number] = issue['title']
+            embeddings_new[number] = embeddings[i]
+            if issue['state'] == 'open':
+                opened[number] = True
+            if issue['state'] == 'closed':
+                closed[number] = True
 
     def embeddings_updated_get(self, repo):
         with self.lock:
@@ -230,97 +282,19 @@ class EmbeddingContext:
             # Consider that if the time hasn't changed, it's the same issue.
             issues = [issue for issue in issues if issue['updated_at'] != date_old]
 
-            numbers_old = data['numbers']
-            titles_old = data['titles']
-            embeddings_old = data['embeddings']
-
-            last_index = len(numbers_old) - 1
-
-            issues = sorted(issues, key=lambda issue: int(issue['number']))
-            issues_clos = [issue for issue in issues if issue['state'] == 'closed']
-            issues_open = [issue for issue in issues if issue['state'] == 'open']
-
-            numbers_clos = [int(issue['number']) for issue in issues_clos]
-            numbers_open = [int(issue['number']) for issue in issues_open]
-
-            old_closed = []
-            for number_clos in numbers_clos:
-                for i_old in range(last_index, -1, -1):
-                    number_old = numbers_old[i_old]
-                    if number_old < number_clos:
-                        break
-                    if number_old == number_clos:
-                        old_closed.append(i_old)
-                        break
-
-            if not old_closed and not issues_open:
-                return data
-
-            mask_open = torch.ones(len(numbers_open), dtype=torch.bool)
-            need_sort = False
-            change_map = []
-            for i_open, number_open in enumerate(numbers_open):
-                for i_old in range(last_index, -1, -1):
-                    number_old = numbers_old[i_old]
-                    if number_old < number_open:
-                        need_sort = need_sort or (i_old != last_index)
-                        break
-                    if number_old == number_open:
-                        change_map.append((i_old, i_open))
-                        mask_open[i_open] = False
-                        break
-
-            if issues_open:
-                texts_to_embed = self.create_strings_to_embbed(issues_open, black_list)
-                embeddings = self.encode(texts_to_embed)
-
-            for i_old, i_open in change_map:
-                titles_old[i_old] = issues_open[i_open]['title']
-                embeddings_old[i_old] = embeddings[i_open]
-
-            if old_closed:
-                total = (len(numbers_old) - len(old_closed)) + (len(numbers_open) - len(change_map))
-                numbers_new = [None] * total
-                titles_new = [None] * total
-                embeddings_new = torch.empty((total, *embeddings_old.shape[1:]), dtype=embeddings_old.dtype, device=embeddings_old.device)
-
-                i_new = 0
-                i_old = 0
-                for i_closed in old_closed + [len(numbers_old)]:
-                    while i_old < i_closed:
-                        numbers_new[i_new] = numbers_old[i_old]
-                        titles_new[i_new] = titles_old[i_old]
-                        embeddings_new[i_new] = embeddings_old[i_old]
-                        i_new += 1
-                        i_old += 1
-                    i_old += 1
-
-                for i_open in range(len(numbers_open)):
-                    if not mask_open[i_open]:
-                        continue
-                    titles_new[i_new] = issues_open[i_open]['title']
-                    numbers_new[i_new] = numbers_open[i_open]
-                    embeddings_new[i_new] = embeddings[i_open]
-                    i_new += 1
-
-                assert i_new == total
-            elif mask_open.any():
-                titles_new = titles_old + [issue['title'] for i, issue in enumerate(issues_open) if mask_open[i]]
-                numbers_new = numbers_old + [number for i, number in enumerate(numbers_open) if mask_open[i]]
-                embeddings_new = torch.cat([embeddings_old, embeddings[mask_open]])
-            else:
-                # Only Updated Data changed
-                return data
+            self.data_ensure_size(repo, int(issues[0]['number']))
 
-            if need_sort:
-                sorted_indices = sorted(range(len(numbers_new)), key=lambda k: numbers_new[k])
-                titles_new = [titles_new[i] for i in sorted_indices]
-                numbers_new = [numbers_new[i] for i in sorted_indices]
-                embeddings_new = embeddings_new[sorted_indices]
+            texts_to_embed = self.create_strings_to_embbed(issues, black_list)
+            embeddings = self.encode(texts_to_embed)
 
-            data['titles'] = titles_new
-            data['numbers'] = numbers_new
-            data['embeddings'] = embeddings_new
+            for i, issue in enumerate(issues):
+                number = int(issue['number'])
+                data['titles'][number] = issue['title']
+                data['embeddings'][number] = embeddings[i]
+                if issue['state'] == 'open':
+                    data['opened'][number] = True
+                if issue['state'] == 'closed':
+                    data['closed'][number] = True
 
 # autopep8: on
         return data
@@ -332,64 +306,55 @@ EMBEDDING_CTX = EmbeddingContext()
 # EMBEDDING_CTX.embeddings_generate('blender', 'blender-addons')
 
 
-def _sort_similarity(data, query_emb, limit):
+# Define your Enum class
+class State(str, Enum):
+    opened = "opened"
+    closed = "closed"
+    all = "all"
+
+
+def _sort_similarity(data: dict,
+                     query_emb: torch.Tensor,
+                     limit: int,
+                     state: State = State.opened) -> list:
     duplicates = []
+    embeddings = data['embeddings']
+    true_indices = None
+
+    if state != State.all:
+        mask = data[state.value]
+        embeddings = embeddings[mask]
+        true_indices = mask.nonzero(as_tuple=True)[0]
+
     ret = util.semantic_search(
-        query_emb, data['embeddings'], top_k=limit, score_function=util.dot_score)
+        query_emb, embeddings, top_k=limit, score_function=util.dot_score)
+
     for score in ret[0]:
         corpus_id = score['corpus_id']
-        text = f"#{data['numbers'][corpus_id]}: {data['titles'][corpus_id]}"
+        number = true_indices[corpus_id].item(
+        ) if true_indices is not None else corpus_id
+        text = f"#{number}: {data['titles'][number]}"
         duplicates.append(text)
 
     return duplicates
 
 
-cached_search = {'text': '', 'repo': '', 'issues': []}
-
-
-def text_search(owner, repo, text_to_embed, limit=None):
-    global cached_search
-    global EMBEDDING_CTX
-    if not text_to_embed:
-        return []
-
-    if text_to_embed == cached_search['text'] and repo == cached_search['repo']:
-        return cached_search['issues'][:limit]
-
-    data = EMBEDDING_CTX.embeddings_updated_get(owner, repo)
-
-    new_embedding = EMBEDDING_CTX.encode([text_to_embed])
-    result = _sort_similarity(data, new_embedding, 500)
-
-    cached_search = {'text': text_to_embed, 'repo': repo, 'issues': result}
-    return result[:limit]
-
+def find_relatedness(repo: str, number: int, limit: int = 20, state: State = State.opened):
+    data = EMBEDDING_CTX.embeddings_updated_get(repo)
 
-def find_relatedness(repo, number, limit=20):
-    gitea_issue = gitea_json_issue_get('blender', repo, number)
-    assert gitea_issue['repository']['owner'] == 'blender'
-    repo = gitea_issue['repository']['name']
-    title = gitea_issue['title']
-    body = gitea_issue['body']
-    number = int(gitea_issue['number'])
+    # Check if the embedding already exists.
+    if data['titles'][number] is not None:
+        new_embedding = data['embeddings'][number]
+    else:
+        gitea_issue = gitea_json_issue_get('blender', repo, number)
+        text_to_embed = _create_issue_string(
+            gitea_issue['title'], gitea_issue['body'])
 
-    data = EMBEDDING_CTX.embeddings_updated_get(repo)
-    new_embedding = None
-
-    # Check if the embedding already exist.
-    for i in range(len(data['numbers']) - 1, -1, -1):
-        number_cached = data['numbers'][i]
-        if number_cached < number:
-            break
-        if number_cached == number:
-            new_embedding = data['embeddings'][i]
-            break
-
-    if new_embedding is None:
-        text_to_embed = _create_issue_string(title, body)
         new_embedding = EMBEDDING_CTX.encode([text_to_embed])
 
-    duplicates = _sort_similarity(data, new_embedding, limit=limit)
+    duplicates = _sort_similarity(
+        data, new_embedding, limit=limit, state=state)
+
     if not duplicates:
         return ''
 
@@ -401,8 +366,8 @@ def find_relatedness(repo, number, limit=20):
 
 
 @router.get("/find_related/{repo}/{number}")
-def find_related(repo: str = 'blender', number: int = 104399, limit: int = 15):
-    related = find_relatedness(repo, number, limit=limit)
+def find_related(repo: str = 'blender', number: int = 104399, limit: int = 15, state: State = State.opened):
+    related = find_relatedness(repo, number, limit=limit, state=state)
     return related
 
 
@@ -425,11 +390,8 @@ if __name__ == "__main__":
             val['embeddings'] = val['embeddings'].to(torch.device('cuda'))
 
         # '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)
-
-        related1 = find_relatedness(issue1, limit=20)
-        related2 = find_relatedness(issue2, limit=20)
+        related1 = find_relatedness('blender', 111434, limit=20)
+        related2 = find_relatedness('blender-addons', 104399, limit=20)
 
         print("These are the 20 most related issues:")
         print(related1)

routers/tool_find_related_cache.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:194833cab1df62e81e34993c5d271c037d16d9353048ed1592f205515f2162ab
-size 21616587
+oid sha256:46c42973a8caaa2f0d4a76ebc6ff16c0b8df927c9b16ba645c3f7155cce84f6a
+size 723382452