Better(?) prompt etc

extract_relations.py

@@ -1,3 +1,4 @@
+import re
 from _arango import arango
 from _llm import LLM
 
@@ -7,7 +8,7 @@ from langchain_text_splitters import CharacterTextSplitter
 # Create an instance of the SentenceSplitter
 text_splitter = CharacterTextSplitter(
     separator="\n\n",
-    chunk_size=8000,
+    chunk_size=6000,
     chunk_overlap=0,
     length_function=len,
     is_separator_regex=False,
@@ -19,30 +20,63 @@ interrogations = [i for i in arango.db.collection("interrogations").all()]
 interrogations = sorted(interrogations, key=lambda x: x["date"])
 
 
+# Sort out interrogations already processed
+all_relations = [i for i in arango.db.collection("all_relations").all()]
+for relation in all_relations:
+    if "interrogation" in relation["context"]:
+        for interrogation in interrogations:
+            if relation["_in"] == interrogation["_id"]:
+                interrogations.remove(interrogation)
+
+
+# Go through all interrogations
 for interrogation in interrogations:
 
     # Get the persons (now updated from the last one, so it should be all persons in the database)
-    persons_docs = [i for i in arango.db.collection("persons").all()]
-    persons = [i["name"] for i in persons_docs]
+    persons_docs = []
+    for collection in ["persons", "other_persons"]:
+        for i in arango.db.collection(collection).all():
+            persons_docs.append(i)
+    persons = [i["name"].strip() for i in persons_docs]
     persons_dict = {i["name"]: i for i in persons_docs}
     persons_string = "\n".join(persons)
 
+    # Get the person who is interrogated
     interrogated_person = interrogation["person"]
     if interrogation["person"] in persons:
         doc = persons_dict[interrogated_person]
         _from = doc["_id"]
+        _from_name = doc["_key"]
     text = interrogation["text"]
+
+    # Make the text into smaller chunks and go through them
     chunks = text_splitter.split_text(text)
+
+    print(len(chunks), 'chunks in', interrogation["person"].upper())
+
+    earlier_mentions = []
+    earlier_mentions_string = ''
+
+    if len(chunks) > 1:
+        part_of_string = ' en del av'
+    else:
+        part_of_string = ''
     for chunk in chunks:
+        if earlier_mentions != []:
+            earlier_mentions_list_string = '\n'.join(list(set(earlier_mentions)))
+            earlier_mentions_string = f"Tidigare i förhöret har personerna nedan nämnts på följande sätt:\n{earlier_mentions_list_string}\n"
+        
         prompt = f''''
-Kolla på texten nedan: \n\
+Texten nedan är{part_of_string} ett polisförhör, kolla på den: \n
+
 TEXT:
 """{chunk}""" \n
-{interrogation["person"]} förhörs. Nämns några personer i listan nedan i själva förhöret? \n
+{interrogation["person"]} förhörs. Jag är intresserad av om någon eller några personer ur listan på personer nedan nämns i själva förhöret.\n
 LISTA PÅ PERSONER:
 {persons_string}\n
-I texten kan en person nämnas med sitt fulla namn, men oftast bara förnamn eller efternamn.
-Svara med fullständiga namn från listan och hur personen nämns i texten på formen "namn;hur personen nämns\n". 
+Nämns någon eller några av personerna i listan i förhöret? Personen kan nämnas med sitt fulla namn, men oftast bara förnamn eller efternamn. Försök alltså förstå om en person som nämns med förnamn är en person i listan.
+{earlier_mentions_string}
+Svara med FULLSTÄNDIGA namn från listan och hur personen nämns i texten på formen "namn;hur personen nämns\n". 
 Nedan är ett exempel för att du ska förstå hur du ska svara:
 
 <EXEMPEL>
@@ -50,8 +84,9 @@ John Lundqvist;John
 Karl Renström; Karl
 </EXEMPEL>
 
-Svara ENBART med personens namn och hur det nämns, på formen "namn;hur personen nämns\n". Svara inte med något resonemang, och enbart med personer som nämns. Om ingen person från listan nämns, svara med None.
-\nPersoner:'''
+Svara ENBART med personens fullständiga namn och hur det nämns. Var noga med att använda formen "namn;hur personen nämns\n". 
+Svara INTE med något resonemang eller något annat, och enbart med personer som nämns i förhöret. Om ingen person från listan nämns, svara med None.
+Personer som nämns:'''
 
         relations = llm.generate(prompt)
 
@@ -60,6 +95,15 @@ Svara ENBART med personens namn och hur det nämns, på formen "namn;hur persone
                 continue
             try:
                 name, mention = relation.split(";", 1)
+                name = name.strip()
+                mention = mention.strip()
+                
+                # Remove numbers at the beginning of the name using regex
+                name = re.sub(r"^\d+\.\s*", "", name)
+                name = re.sub(r"[^a-zA-Z\s-]", "", name)
+
+                original_name = name
+                
             except ValueError:
                 print("\033[91m" + relation + "\033[0m")
                 continue
@@ -67,50 +111,63 @@ Svara ENBART med personens namn och hur det nämns, på formen "namn;hur persone
             if name in persons:
                 doc = persons_dict[name]
                 _to = doc["_id"]
+                _to_name = doc["_key"]
 
+                earlier_mentions.append(f'{name} - {mention}')
             else:
 
                 name_parts = name.split(" ")
-                if ' ' in name and f'{name_parts[1]} {name_parts[0]}' in persons:
-                    doc = persons_dict[f'{name_parts[1]} {name_parts[0]}']
+                part1 = name_parts[0].strip()
+                part2 = name_parts[1].strip() if len(name_parts) > 1 else ""
+                if f'{part2} {part1}' in persons:
+                    doc = persons_dict[f'{part2} {part1}']
                     _to = doc["_id"]
+                    _to_name = doc["_key"]
                 else:
                     arango_doc = {
                         "name": name,
                         "_key": arango.fix_key_name(name),
-                        "interrogated": "Unknown",
                     }
                     # pprint(mention_context)
                     # add = input(f"Add {name} ({mention}) to database? (y/n) >> ")
                     # if add in ["y", ""]:
-                    if arango.fix_key_name(name) not in arango.db.collection("persons"):
-                        arango.db.collection("persons").insert(arango_doc)
-                        doc = arango.db.collection("persons").get(arango_doc["_key"])
+                    if arango.fix_key_name(name) not in arango.db.collection("other_persons"):
+                        doc = arango.db.collection("other_persons").insert(arango_doc)
                         _to = doc["_id"]
+                        _to_name = doc["_key"]
                     else:
-                        doc = arango.db.collection("persons").get(arango_doc["_key"])
+                        doc = arango.db.collection("other_persons").get(arango_doc["_key"])
                         _to = doc["_id"]
-            if _from == _to:
+                        _to_name = doc["_key"]
+            if _from_name == _to_name:
                 continue
 
-            relation_key = arango.fix_key_name(f"{_from}_{_to}__{interrogation['_key']}").replace("persons_", "")
+            relation_key = arango.fix_key_name(f"{_to}__{interrogation['_key']}").replace("persons_", "")
             if arango.db.has_document("all_relations/" + relation_key):
                 continue
             
             # Ask LLM about the context of the mention
-            prompt = f'Nedan är en del av ett förhör med {interrogated_person}.\n\n"""{chunk}"""\n\n{interrogated_person} nämner i förhöret en person vid namn {name}. Exakt vad säger {interrogated_person} om {name}? Svara så kortfattat som möjligt.\n\nSvar:'
+            prompt = f'Nedan är en del av ett förhör med {interrogated_person}.\n\n"""{chunk}"""\n\n{interrogated_person} nämner i förhöret en person vid namn {original_name}. Exakt vad säger {interrogated_person} om {original_name}? Svara så kortfattat som möjligt.\n\nSvar:'
             mention_context = llm.generate(prompt)
+
+            prompt = f'Det här är ett svar från en assistent: """{mention_context}""".\nVerkar {original_name} nämnas? Svara bara med TRUE eller FALSE.\n\nSvar:'
+            is_mentioned = llm.generate(prompt)
+            if 'false' in is_mentioned.lower():
+                print("\033[91m" + f"{original_name} not mentioned" + "\033[0m")
+                print(mention_context, '\n')
+                continue
             arango_doc = {
                     "_key": relation_key,
                     "_from": _from,
                     "_to": _to,
-                    "in": interrogation["_id"],
+                    "_in": interrogation["_id"],
                     "context": "interrogation",
                     "mentioned_as": mention,
                     "mention": mention_context,
+                    "date": interrogation["date"],
                 }
             
-            print("\033[92m" + f'{_from} -> {_to}' + "\033[0m")
+            print("\033[92m" + f'{_from} -> {_to}' + "\033[0m" + f' ({interrogation["_key"]})')
             print(mention_context)
             print()
             arango.db.collection("all_relations").insert(