mala/extract_fup.py

from _llm import LLM
import fitz
from _arango import arango
from openai import OpenAI
from pprint import pprint

class Section:
    def __init__(self, type, page, filename="Förhörsprotokoll.pdf"):
        self.type = type
        self.filename = filename
        self.text = ""
        self.person = ""
        self.role = ""
        self.topic = ""
        self.date = ""
        self.start_page = page
        self.key = ""

    def add_to_arango(self):

        key = arango.fix_key_name(
            f"{self.person}_{self.date}_p.{self.start_page}"
        )
        arango_doc = {
                "_key": key,
                "person": self.person,
                "role": self.role,
                "topic": self.topic,
                "date": self.date,
                "page": self.start_page,
                "text": self.text,
                "filename": self.filename,
            }
        arango.db.collection(self.type).insert(arango_doc, overwrite=True)
        print(f"Added {self.type} to ArangoDB with key {key}")


    def extract_interrogation(self, text):
        self.person = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är förhörd? Svara på formen "Förnamn Efternamn" \n\nFörhörd person:'
        )
        self.role = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är {self.person}? \n\nTitel på förhörd person:'
        )
        self.topic = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vad handlade förhöret om? Svara så kortfattat som möjligt. \n\nFörhörets syfte:'
        )
        self.date = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n När ägde förhöret rum? Svara på formen YYY-MM-DD \n\nFörhörsdatum:'
        )


        self.key = arango.fix_key_name(f"{self.person}_{self.date}_p.{self.start_page}")


    def extract_pm(self, text):
        self.person = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är uppgiftslämnare? Svara på formen "Förnamn Efternamn" \n\nPM:'
        )
        self.role = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vem är {self.person}? Svara "None" om det inte framgår. \n\nTitel på person:'
        )   
        self.topic = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n Vad handlade informationen om? Svara så kortfattat som möjligt. Svara "None" om det inte framgår. \n\Svar:'
        )
        self.date = llm.generate(
            f'Kolla på texten nedan: \n\n """{text}""" \n\n När lämnades informationen? Svara på formen YYY-MM-DD \n\nDatum:'
        )
        self.key = arango.fix_key_name(f"{self.person}_{self.date}_p.{self.start_page}")
        

def new_interrogation(page, section):

    if section.text != "":
        section.add_to_arango()
    section = Section("interrogations", page.number)
    section.extract_interrogation(page.get_text())
    return section


def new_pm(page, section):
    if section.text != "":
        section.add_to_arango()
    section = Section("pms", page.number)
    section.extract_interrogation(page.get_text())
    return section



# * Llama
llm = LLM(chat=False, model="llama3:8b-instruct-q5_K_M")

# Open the PDF file
filename = "Förhörsprotokoll.pdf"
area = fitz.Rect(0, 40, 520, 800)  # To exlude the header

doc = fitz.open(f"pdfs/{filename}")

section = Section("interrogations", 0)


for page in doc.pages(9, len(doc) - 1):
    
    # Get the text from the page
    page_text = page.get_text("text")

    # Check if there is a new interrogation
    control_words_interrogation = [
        "Förhörsdatum",
        "Förhör påbörjat",
        "Förhör avslutat",
        "Förhörssätt",
        "Typ av förhör",
        "Förhörsvittne",
    ]
    n_control_words_interrogation = 0
    for word in control_words_interrogation:
        if word in page_text:
            n_control_words_interrogation += 1
    if n_control_words_interrogation >= 2:
        section = new_interrogation(page, section)
        area = fitz.Rect(0, 400, 520, 800)
    
    else:
        # Check if there is a new PM
        control_words_pm = [
            "PM",
            "Uppgiften avser",
            "Upprättad av",
            "Sätt på vilket uppgift lämnats",
            "Uppgiftslämnare",
        ]
        n_control_words_pm = 0
        for word in control_words_pm:
            if word in page_text:
                n_control_words_pm += 1
        if n_control_words_pm >= 2:
            area = fitz.Rect(0, 400, 520, 800)
            section = new_pm(page, section)
        
        
        else:
            # It's a "normal" page
            area = fitz.Rect(0, 40, 520, 800)  # To exlude the header

    blocks = page.get_text("blocks", clip=area)
    for block in blocks:
        section.text += block[4] + "\n\n"