You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
104 lines
4.2 KiB
104 lines
4.2 KiB
import re |
|
from _arango import arango |
|
from _llm import LLM |
|
from pprint import pprint |
|
from pprint import pprint |
|
from langchain_text_splitters import CharacterTextSplitter |
|
import multiprocessing |
|
from print_color import * |
|
|
|
|
|
def describe_relation(person1, person2, relation, text): |
|
llm = LLM(chat=False, system_prompt="Du ska hitta relationer i en text. Svara alltid enligt angiven form och alltid på svenska.") |
|
prompt = f''' |
|
I texten nedan beskrivs att {person1} och {person2} har relationen "{relation}". Läs texten och sammanfatta kortfattat vad som beskrivs om relationen mellan {person1} och {person2}:\n\n"""{text}"""\n |
|
Svara ENBART med information om relationen, inga hälsningsfraser eller liknande. |
|
Relationen ska vara kortfattad och stämma med texten. Om det inte går att beskriva relationen så svara med "None". |
|
Vad står det om relationen "{relation}" mellan {person1} och {person2}? |
|
''' |
|
response = llm.generate(prompt) |
|
print_rainbow(relation, response) |
|
return response |
|
|
|
|
|
def find_relations(interrogation): |
|
""" |
|
Finds the relations in an interrogation. |
|
|
|
Args: |
|
interrogation (dict): The interrogation. |
|
|
|
Returns: |
|
None |
|
""" |
|
text = interrogation["text"] |
|
text_splitter = CharacterTextSplitter( |
|
separator="\n\n", |
|
chunk_size=6000, |
|
chunk_overlap=0, |
|
length_function=len, |
|
is_separator_regex=False, |
|
) |
|
chunks = text_splitter.split_text(text) |
|
all_relations = [] |
|
llm = LLM(chat=False, system_prompt="Du ska hitta relationer i en text. Svara alltid enligt angiven form och alltid på svenska.") |
|
|
|
for chunk in chunks: |
|
prompt = f"""Nedan är en bit av ett polisförhör med {interrogation['person']}. Jag vill att du hittar alla relationer mellan identifierbara personer som beskrivs i själva förhöret:\n\n{chunk}\n\n |
|
Svara på formen "person1;person2;relation\n". Var noga med hur semikolon används för att skilja personerna och relationen, och ny rad efter varje relation (informationen ska sedan användas för en CSV fil). Svara på svenska. |
|
Nedan är ett påhittat exempel för att du ska förstå hur du kan svara: |
|
<exempel> |
|
person1;person2;gick på grundskolan tillammans, spelade fotboll |
|
person2;person3;gifta sedan 2022 |
|
</exempel> |
|
Beskrivningen av relationen ska vara kortfattad och stämma med texten. |
|
Om det inte finns någon relation, svara "None". |
|
Svara ENBART med relationerna, INGENTING annat som en hälsning eller förklaring. |
|
""" |
|
response = llm.generate(prompt) |
|
print_blue(response) |
|
relations = response.split("\n") |
|
for relation in relations: |
|
if relation == "None" or ';' not in relation: |
|
continue |
|
try: |
|
person1, person2, relation = relation.split(";", 2) |
|
except ValueError as e: |
|
print_red(f"Error: {e}") |
|
print_red(f"Relation: {relation}") |
|
continue |
|
|
|
description = describe_relation(person1, person2, relation, chunk) |
|
|
|
|
|
for r in all_relations: |
|
p1 = r["from"] |
|
p2 = r["to"] |
|
if p1 == person1 and p2 == person2: |
|
r["relations"].append({'relation': relation, 'description': description, 'chunk_number': chunks.index(chunk)}) |
|
break |
|
else: |
|
all_relations.append( |
|
{ |
|
"_key": arango.fix_key_name(f"{person1}-{person2}_{interrogation['_key']}"), |
|
"from": person1, |
|
"to": person2, |
|
"relations": [{'relation': relation, 'description': description, 'chunk_number': chunks.index(chunk)}], |
|
'interrogation': interrogation['_key'], |
|
"chunks": chunks, |
|
} |
|
) |
|
|
|
|
|
return all_relations |
|
|
|
if __name__ == "__main__": |
|
db = arango.db |
|
|
|
q = 'for doc in interrogations return doc' |
|
interrogations = list(db.aql.execute(q)) |
|
for interrogation in interrogations: |
|
relations = find_relations(interrogation) |
|
db.collection('relations').insert_many(relations) |
|
|
|
|
|
|