main.py

import os
import re
import secrets
import string
import yaml
from datetime import datetime
from zipfile import ZipFile

import gradio as gr
import nltk
import pandas as pd
from langchain.embeddings import OpenAIEmbeddings
from langchain.chains import SimpleSequentialChain
from langchain.chat_models import ChatOpenAI
from nltk.tokenize import sent_tokenize
from pandas import DataFrame

import utils
from chains import llm_chains


# download NLTK dependencies
nltk.download("punkt")
nltk.download("stopwords")

# load stop words const.
from nltk.corpus import stopwords
STOP_WORDS = stopwords.words("english")

# load global spacy model
# try:
#     SPACY_MODEL = spacy.load("en_core_web_sm")
# except OSError:
#     print("[spacy] Downloading model: en_core_web_sm")
    
#     spacy.cli.download("en_core_web_sm")
#     SPACY_MODEL = spacy.load("en_core_web_sm")


class Text2KG:
    """Text2KG class."""

    def __init__(self, api_key: str, **kwargs):

        self.llm = ChatOpenAI(openai_api_key=api_key, **kwargs)
        self.embedding = OpenAIEmbeddings(openai_api_key=api_key)

    
    def init(self, steps: list[str]):
        """Initialize Text2KG pipeline from passed steps.
        
        Args:
            *steps (str): Steps to include in pipeline. Must be a top-level name present in
                the schema.yml file
        """
        self.pipeline = SimpleSequentialChain(
            chains=[llm_chains[step](llm=self.llm) for step in steps],
            verbose=False
        )

    
    def run(self, text: str) -> list[dict]:
        """Run Text2KG pipeline on passed text.
        
        Arg:
            text (str): The text input
        
        Returns:
            triplets (list): The list of extracted KG triplets
        """
        triplets = self.pipeline.run(text)

        return triplets


    def clean(self, kg: DataFrame) -> DataFrame:
        """Text2KG post-processing."""
        drop_list = []

        for i, row in kg.iterrows():
            # drop stopwords (e.g. pronouns)
            if (row.subject in STOP_WORDS) or (row.object in STOP_WORDS):
                drop_list.append(i)

            # drop broken triplets
            elif row.hasnans:
                drop_list.append(i)
            
            # lowercase nodes/edges, drop articles
            else:
                article_pattern = r'^(the|a|an) (.+)'
                be_pattern = r'^(are|is) (a )?(.+)'

                kg.at[i, "subject"] = re.sub(article_pattern, r'\2', row.subject.lower())
                kg.at[i, "relation"] = re.sub(be_pattern, r'\3', row.relation.lower())
                kg.at[i, "object"] = re.sub(article_pattern, r'\2', row.object.lower())

        return kg.drop(drop_list)


    def normalize(self, kg: DataFrame, threshold: float=0.3) -> DataFrame:
        """Reduce dimensionality of Text2KG output by merging cosine-similar nodes/edges."""

        ents = pd.concat([kg["subject"], kg["object"]]).unique()
        rels = kg["relation"].unique()

        ent_map = utils.condense_labels(ents, self.embedding.embed_documents, threshold=threshold)
        rel_map = utils.condense_labels(rels, self.embedding.embed_documents, threshold=threshold)

        kg_normal = pd.DataFrame()
        
        kg_normal["subject"] = kg["subject"].map(ent_map)
        kg_normal["relation"] = kg["relation"].map(rel_map)
        kg_normal["object"] = kg["object"].map(ent_map)

        return kg_normal


def extract_knowledge_graph(api_key: str, batch_size: int, modules: list[str], text: str, progress=gr.Progress()):
    """Extract knowledge graph from text.
    
    Args:
        api_key (str): OpenAI API key
        batch_size (int): Number of sentences per forward pass
        modules (list): Additional modules to add before main extraction step
        text (str): Text from which Text2KG will extract knowledge graph from
        progress: Progress bar. The default is gradio's progress bar; for a 
            command line progress bar, set `progress = tqdm`

    Returns:
        zip_path (str): Path to ZIP archive containing outputs
        knowledge_graph (DataFrame): The extracted knowledge graph
    """
    # init
    if api_key == "":
        raise ValueError("API key is required")
    
    pipeline = Text2KG(api_key=api_key, temperature=0.3) # low temp. -> low randomness

    steps = []

    for module in modules:
        m = module.lower().replace(' ', '_')
        steps.append(m)

    if (len(steps) == 0) or (steps[-1] != "triplet_extraction"):
        steps.append("triplet_extraction")

    pipeline.init(steps)

    # split text into batches
    sentences = sent_tokenize(text)
    batches = [" ".join(sentences[i:i+batch_size])
               for i in range(0, len(sentences), batch_size)]
    
    # create KG
    knowledge_graph = []
    
    for i, batch in progress.tqdm(list(enumerate(batches)), 
                                  desc="Processing...", unit="batches"):
        output = pipeline.run(batch)
        [triplet.update({"sentence_id": i}) for triplet in output]

        knowledge_graph.extend(output)


    # convert to df, post-process data
    knowledge_graph = pd.DataFrame(knowledge_graph)
    knowledge_graph = pipeline.clean(knowledge_graph)
    
    # rearrange columns
    knowledge_graph = knowledge_graph[["sentence_id", "subject", "relation", "object"]]

    # metadata
    now = datetime.now()
    date = str(now.date())

    metadata = {
        "_timestamp": now,
        "batch_size": batch_size,
        "modules": steps
    }

    # unique identifier for local saving
    uid = ''.join(secrets.choice(string.ascii_letters)
                  for _ in range(6))
    
    print(f"Run ID: {date}/{uid}")
    
    save_dir = os.path.join(".", "output", date, uid)
    os.makedirs(save_dir, exist_ok=True)


    # save metadata & data
    with open(os.path.join(save_dir, "metadata.yml"), 'w') as f:
        yaml.dump(metadata, f)
    
    batches_df = pd.DataFrame(enumerate(batches), columns=["sentence_id", "text"])
    batches_df.to_csv(os.path.join(save_dir, "sentences.txt"), 
                     index=False)

    knowledge_graph.to_csv(os.path.join(save_dir, "kg.txt"), 
                           index=False)    
    

    # create ZIP file
    zip_path = os.path.join(save_dir, "output.zip")

    with ZipFile(zip_path, 'w') as zipObj:

        zipObj.write(os.path.join(save_dir, "metadata.yml"))
        zipObj.write(os.path.join(save_dir, "sentences.txt"))
        zipObj.write(os.path.join(save_dir, "kg.txt"))

    return zip_path, knowledge_graph


class App:
    def __init__(self):
        demo = gr.Interface(
            fn=extract_knowledge_graph,
            title="Text2KG",
            inputs=[
                gr.Textbox(placeholder="API key...", label="OpenAI API Key", type="password"),
                gr.Slider(minimum=1, maximum=10, step=1, label="Sentence Batch Size"),
                gr.CheckboxGroup(choices=["Clause Deconstruction"], label="Optional Modules"),
                gr.Textbox(lines=2, placeholder="Text Here...", label="Input Text"),
            ],
            outputs=[
                gr.File(label="Knowledge Graph"),
                gr.DataFrame(label="Preview", 
                             headers=["sentence_id", "subject", "relation", "object"], 
                             max_rows=10, 
                             overflow_row_behaviour="paginate")
            ],
            examples=[[None, 1, [], ("All cells share four common components: "
                                        "1) a plasma membrane, an outer covering that separates the "
                                        "cell's interior from its surrounding environment; 2) cytoplasm, "
                                        "consisting of a jelly-like cytosol within the cell in which "
                                        "there are other cellular components; 3) DNA, the cell's genetic "
                                        "material; and 4) ribosomes, which synthesize proteins. However, "
                                        "prokaryotes differ from eukaryotic cells in several ways. A "
                                        "prokaryote is a simple, mostly single-celled (unicellular) "
                                        "organism that lacks a nucleus, or any other membrane-bound "
                                        "organelle. We will shortly come to see that this is significantly "
                                        "different in eukaryotes. Prokaryotic DNA is in the cell's central "
                                        "part: the nucleoid.")]],
            allow_flagging="never",
            cache_examples=False
        )
        demo.queue().launch(share=False)


if __name__ == "__main__":
    App()