Lasse Server
2 years ago
parent
fe972973c4
commit
f3ee5ba669
2 changed files with 176 additions and 176 deletions
@ -0,0 +1,176 @@ |
||||
import re |
||||
from _arango import arango |
||||
from _llm import LLM |
||||
|
||||
from pprint import pprint |
||||
from langchain_text_splitters import CharacterTextSplitter |
||||
|
||||
# Create an instance of the SentenceSplitter |
||||
text_splitter = CharacterTextSplitter( |
||||
separator="\n\n", |
||||
chunk_size=6000, |
||||
chunk_overlap=0, |
||||
length_function=len, |
||||
is_separator_regex=False, |
||||
) |
||||
|
||||
llm = LLM(chat=False) |
||||
|
||||
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 = [] |
||||
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'''' |
||||
Texten nedan är{part_of_string} ett polisförhör, kolla på den: \n |
||||
|
||||
TEXT: |
||||
"""{chunk}""" \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 |
||||
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> |
||||
John Lundqvist;John |
||||
Karl Renström; Karl |
||||
</EXEMPEL> |
||||
|
||||
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) |
||||
|
||||
for relation in relations.replace('*', '').split("\n"): |
||||
if relation == "None": |
||||
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(" ") |
||||
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), |
||||
} |
||||
# 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("other_persons"): |
||||
doc = arango.db.collection("other_persons").insert(arango_doc) |
||||
_to = doc["_id"] |
||||
_to_name = doc["_key"] |
||||
else: |
||||
doc = arango.db.collection("other_persons").get(arango_doc["_key"]) |
||||
_to = doc["_id"] |
||||
_to_name = doc["_key"] |
||||
if _from_name == _to_name: |
||||
continue |
||||
|
||||
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 {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"], |
||||
"context": "interrogation", |
||||
"mentioned_as": mention, |
||||
"mention": mention_context, |
||||
"date": interrogation["date"], |
||||
} |
||||
|
||||
print("\033[92m" + f'{_from} -> {_to}' + "\033[0m" + f' ({interrogation["_key"]})') |
||||
print(mention_context) |
||||
print() |
||||
arango.db.collection("all_relations").insert( |
||||
arango_doc, |
||||
overwrite=True, |
||||
) |
||||
@ -1,176 +0,0 @@ |
||||
import re |
||||
from _arango import arango |
||||
from _llm import LLM |
||||
|
||||
from pprint import pprint |
||||
from langchain_text_splitters import CharacterTextSplitter |
||||
|
||||
# Create an instance of the SentenceSplitter |
||||
text_splitter = CharacterTextSplitter( |
||||
separator="\n\n", |
||||
chunk_size=6000, |
||||
chunk_overlap=0, |
||||
length_function=len, |
||||
is_separator_regex=False, |
||||
) |
||||
|
||||
llm = LLM(chat=False) |
||||
|
||||
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 = [] |
||||
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'''' |
||||
Texten nedan är{part_of_string} ett polisförhör, kolla på den: \n |
||||
|
||||
TEXT: |
||||
"""{chunk}""" \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 |
||||
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> |
||||
John Lundqvist;John |
||||
Karl Renström; Karl |
||||
</EXEMPEL> |
||||
|
||||
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) |
||||
|
||||
for relation in relations.replace('*', '').split("\n"): |
||||
if relation == "None": |
||||
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(" ") |
||||
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), |
||||
} |
||||
# 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("other_persons"): |
||||
doc = arango.db.collection("other_persons").insert(arango_doc) |
||||
_to = doc["_id"] |
||||
_to_name = doc["_key"] |
||||
else: |
||||
doc = arango.db.collection("other_persons").get(arango_doc["_key"]) |
||||
_to = doc["_id"] |
||||
_to_name = doc["_key"] |
||||
if _from_name == _to_name: |
||||
continue |
||||
|
||||
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 {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"], |
||||
"context": "interrogation", |
||||
"mentioned_as": mention, |
||||
"mention": mention_context, |
||||
"date": interrogation["date"], |
||||
} |
||||
|
||||
print("\033[92m" + f'{_from} -> {_to}' + "\033[0m" + f' ({interrogation["_key"]})') |
||||
print(mention_context) |
||||
print() |
||||
arango.db.collection("all_relations").insert( |
||||
arango_doc, |
||||
overwrite=True, |
||||
) |
||||
Loading…
Reference in new issue