Merge branch 'main' of https://huggingface.co/spaces/jhatchett/Text2KG

main.py

@@ -3,6 +3,7 @@ import re
 import secrets
 import string
 import yaml
+import yaml
 from datetime import datetime
 from zipfile import ZipFile
 
@@ -10,6 +11,7 @@ import gradio as gr
 import nltk
 import pandas as pd
 from langchain.embeddings import OpenAIEmbeddings
+from langchain.embeddings import OpenAIEmbeddings
 from langchain.chains import SimpleSequentialChain
 from langchain.chat_models import ChatOpenAI
 from nltk.tokenize import sent_tokenize
@@ -139,6 +141,14 @@ def extract_knowledge_graph(api_key: str, batch_size: int, modules: list[str], t
 
     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")
+    steps = []
+
     for module in modules:
         m = module.lower().replace(' ', '_')
         steps.append(m)
@@ -148,10 +158,13 @@ def extract_knowledge_graph(api_key: str, batch_size: int, modules: list[str], t
 
     pipeline.init(steps)
 
+    # split text into batches
     # split text into batches
     sentences = sent_tokenize(text)
     batches = [" ".join(sentences[i:i+batch_size])
                for i in range(0, len(sentences), batch_size)]
+    batches = [" ".join(sentences[i:i+batch_size])
+               for i in range(0, len(sentences), batch_size)]
     
     # create KG
     knowledge_graph = []
@@ -179,6 +192,9 @@ def extract_knowledge_graph(api_key: str, batch_size: int, modules: list[str], t
         "_timestamp": now,
         "batch_size": batch_size,
         "modules": steps
+        "_timestamp": now,
+        "batch_size": batch_size,
+        "modules": steps
     }
 
     # unique identifier for local saving

requirements.txt

@@ -1,8 +1,8 @@
 gradio==4.2.0
 langchain==0.0.335
-matplotlib==3.7.1
 nltk==3.7
 openai==0.27.4
 pandas==2.0.3
 PyYAML==6.0
+scikit-learn==1.2.2
 tqdm==4.65.0