Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| import secrets | |
| import string | |
| from datetime import datetime | |
| from zipfile import ZipFile | |
| import gradio as gr | |
| import pandas as pd | |
| from langchain.chains import SimpleSequentialChain | |
| from langchain.chat_models import ChatOpenAI | |
| from nltk.tokenize import sent_tokenize | |
| import utils | |
| from chains import chains | |
| class Text2KG: | |
| """Text2KG class.""" | |
| def __init__(self, api_key: str, **kwargs): | |
| self.model = ChatOpenAI(openai_api_key=api_key, **kwargs) | |
| def init_pipeline(self, *steps: 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=[chains[step](llm=self.model) 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 create_knowledge_graph(api_key: str, ngram_size: int, axiomatize: bool, text: str, progress=gr.Progress()): | |
| """Create 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 | |
| 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 | |
| """ | |
| # init | |
| if api_key == "": | |
| raise ValueError("API key is required") | |
| model = Text2KG(api_key=api_key, temperature=0.3) # low temp. -> low randomness | |
| if axiomatize: | |
| steps = ["text2axiom", "extract_triplets"] | |
| else: | |
| steps = ["extract_triplets"] | |
| model.init_pipeline(*steps) | |
| # split text into ngrams | |
| sentences = sent_tokenize(text) | |
| ngrams = [" ".join(sentences[i:i+ngram_size]) | |
| for i in range(0, len(sentences), ngram_size)] | |
| # create KG | |
| knowledge_graph = [] | |
| for i, ngram in progress.tqdm(enumerate(ngrams), desc="Processing...", total=len(ngrams)): | |
| output = model.run(ngram) | |
| [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 = utils.process(knowledge_graph) | |
| # rearrange columns | |
| knowledge_graph = knowledge_graph[["sentence_id", "subject", "relation", "object"]] | |
| # 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 | |
| } | |
| # unique identifier for saving | |
| uid = ''.join(secrets.choice(string.ascii_letters) | |
| for _ in range(6)) | |
| 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) | |
| ngrams_df = pd.DataFrame(enumerate(ngrams), columns=["sentence_id", "text"]) | |
| ngrams_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.json")) | |
| zipObj.write(os.path.join(save_dir, "sentences.txt")) | |
| zipObj.write(os.path.join(save_dir, "kg.txt")) | |
| return knowledge_graph, zip_path | |
| 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, | |
| 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.Textbox(lines=2, placeholder="Text Here...", label="Input Text"), | |
| ], | |
| outputs=[ | |
| gr.DataFrame(label="Knowledge Graph Triplets", | |
| headers=["sentence_id", "subject", "relation", "object"], | |
| max_rows=10, | |
| overflow_row_behaviour="show_ends"), | |
| gr.File(label="Knowledge Graph") | |
| ], | |
| examples=[["", 1, False, ("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.launch(share=False) | |
| if __name__ == "__main__": | |
| App() |