|
|
import multiprocessing |
|
|
from _llm import LLM |
|
|
from _arango import arango |
|
|
from langchain_text_splitters import CharacterTextSplitter |
|
|
import difflib |
|
|
import re |
|
|
import random |
|
|
from time import sleep |
|
|
import traceback |
|
|
from pprint import pprint |
|
|
|
|
|
def check_name(name, persons): |
|
|
valid_name = True |
|
|
name_parts = name.split() |
|
|
if len(name_parts) > 1: |
|
|
if name_parts[0].istitle() and name_parts[1].istitle(): |
|
|
valid_name = False |
|
|
return valid_name |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def execute_query_with_retry(db, query, max_retries=5, delay=2): |
|
|
for attempt in range(max_retries): |
|
|
try: |
|
|
result = db.aql.execute(query) |
|
|
# If the result is a string, raise an exception |
|
|
if isinstance(result, str): |
|
|
raise ValueError(f"Unexpected result from database: {result}") |
|
|
return list(result) |
|
|
except Exception as e: |
|
|
print(f"Error executing query, attempt {attempt+1}: {e}") |
|
|
sleep(delay) |
|
|
# If we've exhausted all retries, re-raise the last exception |
|
|
raise |
|
|
|
|
|
# Then, in your extract_persons function: |
|
|
def extract_persons(interrogation): |
|
|
print(interrogation["_key"].upper()) |
|
|
q = "for doc in persons filter doc.other != true return doc" |
|
|
result = execute_query_with_retry(db, q) |
|
|
|
|
|
persons_docs = list(result) |
|
|
|
|
|
persons = [i["name"].strip() for i in persons_docs] |
|
|
|
|
|
persons_dict = {i["name"]: i for i in persons_docs} |
|
|
|
|
|
text_splitter = CharacterTextSplitter( |
|
|
separator="\n\n", |
|
|
chunk_size=4000, |
|
|
chunk_overlap=0, |
|
|
length_function=len, |
|
|
is_separator_regex=False, |
|
|
) |
|
|
chunks = text_splitter.split_text(interrogation["text"]) |
|
|
llm = LLM( |
|
|
chat=True, |
|
|
system_prompt="Du är en assistent som hjälper till att hitta personer i ett polisförhör. Du får en del av texten från förhöret åt gången. Svara bara när personen finns i den del du får, hitta inte på personer.", |
|
|
) |
|
|
names = [] |
|
|
for chunk in chunks: |
|
|
# Find persons in the text |
|
|
prompt = f'''Det här är en text från ett polisförhör där {interrogation["person"]} förhörs:\n |
|
|
"""{chunk}"""\n |
|
|
Vilka personer nämns i texten som inte förekommit tidigare? Svara ENBART med en pythonformaterad lista av namn. |
|
|
Exempel på svar för att du ska förstå formen: "["namn1", "namn2", "namn3"]". |
|
|
Jag är inte intresserad av förhörsledaren eller personen som förhörs.''' |
|
|
response = llm.generate(prompt) |
|
|
response = re.sub(r"[^a-zA-ZåäöÅÄÖ\- ,]", "", response).replace(" namn ", "") |
|
|
print(response) |
|
|
for name in [i.strip() for i in response.split(",") if len(i) > 2]: |
|
|
if name not in names: |
|
|
names.append(name) |
|
|
|
|
|
for name in names: |
|
|
try: |
|
|
# Compare the person to a list of known persons |
|
|
prompt = f'''Jag vill veta vem {name} är. Kolla på förhöret nedan och svara om du hittar något om personen där. |
|
|
"""{chunk}"""\n |
|
|
Vem är {name}? Svara bara med sådant som finns i texten.''' |
|
|
info = llm.generate(prompt) |
|
|
person = None |
|
|
# Reverse name |
|
|
if name in persons: |
|
|
person = persons_dict[name] |
|
|
elif name.split().reverse() in persons: |
|
|
print("Vände och hittade ✌️", name.split().reverse()) |
|
|
person = persons_dict[name.split().reverse()] |
|
|
else: |
|
|
closest_matchs = difflib.get_close_matches(name, persons, n=8) |
|
|
persons_string = "\n".join(closest_matchs) |
|
|
prompt = f"""Jag behöver identifiera {name}. Nedan är en lista på personer det kanske skulle kunna vara:\n |
|
|
{persons_string}\n |
|
|
Är {name} någon av dessa personer? Ofta står personen bara med sitt förnamn eller efternamn, men försök att lista ut om det är någon av personerna ovan. Namnet kan också vara felstavat, men inte ett helt annat namn. |
|
|
Svara BARA med namnet på personen ur listan, och bara om du är helst säker på att det verkligen är samma person. Är du inte säker så svara "None".""" |
|
|
answer_person = llm.generate(prompt) |
|
|
|
|
|
if answer_person in persons and check_name(name, persons): |
|
|
person = persons_dict[answer_person] |
|
|
else: |
|
|
q = f"for doc in persons return {{'name': doc['name'], 'info': doc['info']}}" |
|
|
persons_arango_docs = list(db.aql.execute(q)) |
|
|
|
|
|
persons_with_info = [ |
|
|
f"{name} - {info}" |
|
|
for i in persons_arango_docs |
|
|
for name, info in i.items() |
|
|
] |
|
|
persons_with_info_string = "\n".join(persons_with_info) |
|
|
prompt = f"""Här är mer information om möjliga personer:\n |
|
|
{persons_with_info_string}\n |
|
|
|
|
|
Försök att utifrån informationen ovan samt förhöret du tagit del av i tidigare meddelanden identifiera {name}. |
|
|
Svara BARA med namnet på personen ur listan, och bara om du är helst säker på att det verkligen är samma person. Är du inte säker så svara "None".""" |
|
|
|
|
|
answer_person = llm.generate(prompt) |
|
|
|
|
|
if answer_person in persons and check_name(name, persons): |
|
|
person = persons_dict[answer_person] |
|
|
|
|
|
if person: |
|
|
print(f"\033[92m{name} found in database: {person['name']}\033[0m") |
|
|
|
|
|
if "info" not in person: |
|
|
person["info"] = [] |
|
|
if info not in person["info"]: |
|
|
person["info"].append(info) |
|
|
if "mentioned_in_interrogation" not in person: |
|
|
person["mentioned_in_interrogation"] = [] |
|
|
if interrogation["_key"] not in person["mentioned_in_interrogation"]: |
|
|
person["mentioned_in_interrogation"].append(interrogation["_key"]) |
|
|
db.collection("persons").update(person, check_rev=False) |
|
|
db.collection("all_relations").insert( |
|
|
{ |
|
|
"_from": interrogation["person_id"], |
|
|
"_to": person["_id"], |
|
|
"relation": "mentioned_by", |
|
|
"mentions": [{'interrogation': interrogation["_key"], "date": interrogation["date"], "mentioned_as": info}], |
|
|
"_key": f'{interrogation["_key"]}_{person["_key"]}' |
|
|
}, |
|
|
overwrite_mode="update", |
|
|
merge=True, |
|
|
) |
|
|
else: |
|
|
print(f"\033[91m{name} not found in database\033[0m") |
|
|
|
|
|
_key = arango.fix_key_name(name) |
|
|
doc = db.collection("persons").insert( |
|
|
{ |
|
|
"_key": _key, |
|
|
"name": name, |
|
|
"info": info, |
|
|
"other": True, |
|
|
"mentioned_in_interrogation": [interrogation["_key"]], |
|
|
}, |
|
|
overwrite_mode="update", |
|
|
merge=True, |
|
|
) |
|
|
|
|
|
db.collection("all_relations").insert( |
|
|
{ |
|
|
"_from": interrogation["person_id"], |
|
|
"_to": doc["_id"], |
|
|
"relation": "mentioned_by", |
|
|
'other': True, |
|
|
"mentions": [{'interrogation': interrogation["_key"], "date": interrogation["date"], "mentioned_as": info}], |
|
|
"_key": f'{interrogation["_key"]}_{doc["_key"]}' |
|
|
}, |
|
|
overwrite_mode="update", |
|
|
merge=True, |
|
|
) |
|
|
except Exception as e: |
|
|
traceback.print_exc() |
|
|
print(f"\033[91mError when processing {name}: {e}\033[0m") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
db = arango.db |
|
|
q = "for doc in interrogations return doc" |
|
|
interrogations = list(db.aql.execute(q)) |
|
|
interrogations.sort(key=lambda x: x["date"]) |
|
|
|
|
|
persons = list(db.collection("persons").all()) |
|
|
|
|
|
interrogations_done = [] |
|
|
for person in persons: |
|
|
if "mentioned_in_interrogation" in person and person["mentioned_in_interrogation"]: |
|
|
for interrogation in person["mentioned_in_interrogation"]: |
|
|
interrogations_done.append(interrogation) |
|
|
|
|
|
|
|
|
interrogations = [interrogation for interrogation in interrogations if interrogation["_key"] not in set(interrogations_done)] |
|
|
print("Number of interrogations to process:", len(interrogations)) |
|
|
|
|
|
# q = 'for doc in interrogations return doc' |
|
|
# interrogations = list(db.aql.execute(q)) |
|
|
|
|
|
# # Filter out interrogations that have their _key in the rumors collection |
|
|
# q = 'for rumor in rumors return rumor._key' |
|
|
# rumors = list(db.aql.execute(q)) |
|
|
# interrogations = [interrogation for interrogation in interrogations if interrogation['_key'] not in rumors] |
|
|
# print('Number of interrogations to process:', len(interrogations)) |
|
|
|
|
|
# for interrogation in interrogations: |
|
|
# extract_persons(interrogation) |
|
|
with multiprocessing.Pool() as pool: |
|
|
pool.map(extract_persons, interrogations)
|
|
|
|
