Update pipeline

README.md

@@ -1,6 +1,6 @@
 # Text2KG
 
-Using large language models (ChatGPT) to automatically construct a knowledge graph from unstructured plain text.
+We introduce Text2KG – an intuitive, domain-independent tool that leverages the creative generative ability of GPT-3.5 in the KG construction process. Text2KG automates and accelerates the construction of KGs from unstructured plain text, reducing the need for traditionally-used human labor and computer resources. Our approach incorporates a novel, clause-based text simplification step, reducing the processing of even the most extensive corpora down to the order of minutes. With Text2KG, we aim to streamline the creation of databases from natural language, offering a robust, cost-effective, and user-friendly solution for KG construction.
 
 ## Usage
 
@@ -25,21 +25,24 @@ in the repository's directory.
 Import the primary pipeline method using
 
 ```python
->>> from main import create_knowledge_graph
+>>> from main import extract_knowledge_graph
 ```
 
-**`create_knowledge_graph` parameters**
+**`extract_knowledge_graph` parameters**
 
 ```
 api_key (str)
     OpenAI API key
-ngram_size (int)
+
+batch_size (int)
     Number of sentences per forward pass
-axiomatize (bool)
-    Whether to decompose sentences into simpler axioms as a
-    pre-processing step. Doubles the amount of calls to ChatGPT
+
+modules (list)
+    Additional modules to add before main extraction process (triplet_extraction). Must be a valid name in schema.yml
+
 text (str)
     Input text to extract knowledge graph from
+
 progress
     Progress bar. The default is Gradio's progress bar; 
     set `progress = tqdm` for implementations outside of Gradio
@@ -53,10 +56,10 @@ Read more [here](https://www.gradio.app/docs/python-client).
 
 ```
 chains.py
-    Converts schema.yml items to LangChain chains
+    Converts the items in schema.yml to LangChain modules
 
-environment.yml
-    Contains packages required to run environment
+requirements.txt
+    Contains packages required to run Text2KG
 
 main.py
     Main pipeline/app code
@@ -65,7 +68,7 @@ README.md
     This file
 
 schema.yml
-    Contains definitions of prompts
+    Contains definitions of modules -- prompts + output parsers
 
 utils.py
     Contains helper functions

chains.py

@@ -10,7 +10,7 @@ with open("./schema.yml") as f:
     schema = yaml.safe_load(f)
 
 
-class AxiomParser(NumberedListOutputParser):
+class ClauseParser(NumberedListOutputParser):
 
     def parse(self, text: str) -> str:
         axioms = super().parse(text=text)
@@ -34,13 +34,13 @@ class TripletParser(NumberedListOutputParser):
         return super().get_format_instructions()
 
 
-chains = {}
+llm_chains = {}
 
 for scheme in schema:
     parser = schema[scheme]["parser"]
     prompts = schema[scheme]["prompts"]
 
-    chains[scheme] = partial(
+    llm_chains[scheme] = partial(
         LLMChain, 
         output_parser=eval(f'{parser}()'), 
         prompt=ChatPromptTemplate.from_messages(list(prompts.items()))

main.py

@@ -1,18 +1,40 @@
-import json
 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 chains
+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:
@@ -20,10 +42,11 @@ class Text2KG:
 
     def __init__(self, api_key: str, **kwargs):
 
-        self.model = ChatOpenAI(openai_api_key=api_key, **kwargs)
+        self.llm = ChatOpenAI(openai_api_key=api_key, **kwargs)
+        self.embedding = OpenAIEmbeddings(openai_api_key=api_key)
 
     
-    def init_pipeline(self, *steps: str):
+    def init(self, steps: list[str]):
         """Initialize Text2KG pipeline from passed steps.
         
         Args:
@@ -31,7 +54,7 @@ class Text2KG:
                 the schema.yml file
         """
         self.pipeline = SimpleSequentialChain(
-            chains=[chains[step](llm=self.model) for step in steps],
+            chains=[llm_chains[step](llm=self.llm) for step in steps],
             verbose=False
         )
 
@@ -50,45 +73,92 @@ class Text2KG:
         return triplets
 
 
-def create_knowledge_graph(api_key: str, ngram_size: int, axiomatize: bool, text: str, progress=gr.Progress()):
-    """Create knowledge graph from text.
+    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
-        ngram_size (int): Number of sentences per forward pass
-        axiomatize (bool): Whether to decompose sentences into simpler axioms as 
-            a pre-processing step. Doubles the amount of calls to ChatGPT
+        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:
-        knowledge_graph (DataFrame): The extracted knowledge graph
         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")
     
-    model = Text2KG(api_key=api_key, temperature=0.3) # low temp. -> low randomness
+    pipeline = Text2KG(api_key=api_key, temperature=0.3) # low temp. -> low randomness
+
+    steps = []
 
-    if axiomatize:
-        steps  = ["text2axiom", "extract_triplets"]
-    else:
-        steps = ["extract_triplets"]
+    for module in modules:
+        m = module.lower().replace(' ', '_')
+        steps.append(m)
 
-    model.init_pipeline(*steps)
+    if (len(steps) == 0) or (steps[-1] != "triplet_extraction"):
+        steps.append("triplet_extraction")
 
-    # split text into ngrams
+    pipeline.init(steps)
+
+    # split text into batches
     sentences = sent_tokenize(text)
-    ngrams = [" ".join(sentences[i:i+ngram_size]) 
-              for i in range(0, len(sentences), ngram_size)]
+    batches = [" ".join(sentences[i:i+batch_size])
+               for i in range(0, len(sentences), batch_size)]
     
     # create KG
     knowledge_graph = []
     
-    for i, ngram in progress.tqdm(enumerate(ngrams), desc="Processing...", total=len(ngrams)):
-        output = model.run(ngram)
+    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)
@@ -96,7 +166,7 @@ def create_knowledge_graph(api_key: str, ngram_size: int, axiomatize: bool, text
 
     # convert to df, post-process data
     knowledge_graph = pd.DataFrame(knowledge_graph)
-    knowledge_graph = utils.process(knowledge_graph)
+    knowledge_graph = pipeline.clean(knowledge_graph)
     
     # rearrange columns
     knowledge_graph = knowledge_graph[["sentence_id", "subject", "relation", "object"]]
@@ -104,28 +174,29 @@ def create_knowledge_graph(api_key: str, ngram_size: int, axiomatize: bool, text
     # metadata
     now = datetime.now()
     date = str(now.date())
-    timestamp = now.strftime("%Y%m%d%H%M%S")
 
     metadata = {
-        "timestamp": timestamp,
-        "batch_size": ngram_size,
-        "axiom_decomposition": axiomatize
+        "_timestamp": now,
+        "batch_size": batch_size,
+        "modules": steps
     }
 
-    # unique identifier for saving
+    # 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.json"), 'w') as f:
-        json.dump(metadata, f)
+    with open(os.path.join(save_dir, "metadata.yml"), 'w') as f:
+        yaml.dump(metadata, f)
     
-    ngrams_df = pd.DataFrame(enumerate(ngrams), columns=["sentence_id", "text"])
-    ngrams_df.to_csv(os.path.join(save_dir, "sentences.txt"), 
+    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"), 
@@ -137,37 +208,32 @@ def create_knowledge_graph(api_key: str, ngram_size: int, axiomatize: bool, text
 
     with ZipFile(zip_path, 'w') as zipObj:
 
-        zipObj.write(os.path.join(save_dir, "metadata.json"))
+        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 knowledge_graph, zip_path
+    return zip_path, knowledge_graph
 
 
 class App:
     def __init__(self):
-        description = (
-            "# Text2KG\n\n"
-            "Text2KG is a framework that uses ChatGPT to automatically creates knowledge graphs from plain text.\n\n"
-            "**Usage:** (1) configure the pipeline; (2) add the text that will be processed"
-        )
         demo = gr.Interface(
-            fn=create_knowledge_graph,
-            description=description,
+            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", info="Number of sentences per forward pass? Affects the number of calls made to ChatGPT.", ),
-                gr.Checkbox(label="Axiom Decomposition", info="Decompose sentences into simpler axioms? (ex: \"I like cats and dogs.\" = \"I like cats. I like dogs.\")\n\nDoubles the number of calls to ChatGPT."),
+                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.DataFrame(label="Knowledge Graph Triplets", 
+                gr.File(label="Knowledge Graph"),
+                gr.DataFrame(label="Preview", 
                              headers=["sentence_id", "subject", "relation", "object"], 
                              max_rows=10, 
-                             overflow_row_behaviour="show_ends"),
-                gr.File(label="Knowledge Graph")
+                             overflow_row_behaviour="paginate")
             ],
-            examples=[["", 1, False, ("All cells share four common components: "
+            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 "
@@ -182,7 +248,7 @@ class App:
             allow_flagging="never",
             cache_examples=False
         )
-        demo.launch(share=False)
+        demo.queue().launch(share=False)
 
 
 if __name__ == "__main__":

schema.yml

@@ -1,5 +1,5 @@
-text2axiom:
-  parser: AxiomParser
+clause_deconstruction:
+  parser: ClauseParser
   prompts:
     system: |
       You are a sentence parsing agent helping to construct a knowledge graph.
@@ -18,7 +18,7 @@ text2axiom:
     human: |
       {text}
 
-extract_triplets:
+triplet_extraction:
   parser: TripletParser
   prompts:
     system: |

utils.py

@@ -1,34 +1,31 @@
-import nltk
-
-nltk.download("punkt")
-nltk.download("stopwords")
-
-import re
-from nltk.corpus import stopwords
-from pandas import DataFrame
-
-stop_words = stopwords.words("english")
-
-def process(df: DataFrame):
-    """Text2KG post-processing."""
-    drop_list = []
-
-    for i, row in df.iterrows():
-        # remove stopwords (pronouns)
-        if (row.subject in stop_words) or (row.object in stop_words):
-            drop_list.append(i)
-        
-        # remove broken triplets
-        elif row.hasnans:
-            drop_list.append(i)
-        
-        # lowercase nodes/edges, remove articles
-        else:
-            article_pattern = r'^(the|a|an) (.+)'
-            be_pattern = r'^(are|is) (a )?(.+)'
-
-            df.at[i, "subject"] = re.sub(article_pattern, r'\2', row.subject.lower())
-            df.at[i, "relation"] = re.sub(be_pattern, r'\3', row.relation.lower())
-            df.at[i, "object"] = re.sub(article_pattern, r'\2', row.object.lower())
-
-    return df.drop(drop_list)
+import numpy as np
+from typing import Callable
+
+from sklearn.cluster import AgglomerativeClustering
+
+
+def condense_labels(labels: np.ndarray, embedding_func: Callable, threshold: float=0.5):
+    """Combine cosine-similar labels under same name."""
+
+    embeddings = np.array(embedding_func(labels))
+    
+    clustering = AgglomerativeClustering(
+        n_clusters=None, 
+        distance_threshold=threshold
+    ).fit(embeddings)
+
+    clusters = [np.where(clustering.labels_ == l)[0] 
+                for l in range(clustering.n_clusters_)]
+
+    clusters_reduced = []
+    
+    for c in clusters:
+        embs = embeddings[c]
+        centroid = np.mean(embs)
+
+        idx = c[np.argmin(np.linalg.norm(embs - centroid, axis=1))]
+        clusters_reduced.append(idx)
+
+    old2new = {old_id: new_id for old_ids, new_id in zip(clusters, clusters_reduced) for old_id in old_ids}
+    
+    return {labels[i]: labels[j] for i, j in old2new.items()}