diff --git a/extract_relations.py b/extract_relations.py index 923203d..e6b1b74 100644 --- a/extract_relations.py +++ b/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: @@ -50,8 +84,9 @@ John Lundqvist;John Karl Renström; Karl -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,57 +95,79 @@ 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 - + 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(