Spaces:
Sleeping
Sleeping
| import time | |
| from my_lib.chatbot import ChatGPT | |
| import re | |
| class BasePipeline: | |
| chat = ChatGPT(model="gpt-3.5-turbo") | |
| def add_task(self, task, **inputs): | |
| args = task["args"] | |
| args_supplied = {arg: arg in inputs for arg in args} | |
| if all(args_supplied.values()): | |
| formatted_task = f"[TASK]\n{task['description']}" | |
| if len(args): | |
| formatted_task += f"\n\n[INPUT]\n" + '\n'.join([f'{arg}: {inputs[arg]}' for arg in args]) | |
| self.pipe.append(formatted_task) | |
| else: | |
| raise ValueError( | |
| f"Missing required argument(s): {[arg for arg in args if not args_supplied[arg]]}" | |
| ) | |
| def compile_tasks(self, inputs_by_idx): | |
| self.pipe = [] | |
| for idx, task in enumerate(self.prompt): | |
| self.add_task(task, **inputs_by_idx.get(idx, {})) | |
| def forward(self): | |
| self.history = [] | |
| for task in self.pipe: | |
| output = self.chat(task) | |
| time.sleep(0.2) | |
| self.history.append(self.chat.history) | |
| self.chat.clear_history() | |
| return self.postprocess(output) | |
| def postprocess(self): | |
| raise NotImplementedError | |
| class Text2KG(BasePipeline): | |
| def __init__(self, recipe: dict): | |
| self.name = recipe["name"] | |
| self.prompt = recipe["prompt"] | |
| self.chat = ChatGPT( | |
| model="gpt-3.5-turbo", | |
| init=( | |
| "You are a sentence parsing agent helping to construct a knowledge graph." | |
| ), | |
| temperature=0.3 | |
| ) | |
| def __call__(self, text): | |
| self.compile_tasks( | |
| inputs_by_idx={0: {"text": text}} | |
| ) | |
| return self.forward() | |
| def __repr__(self): | |
| return f"Text2KG(recipe={self.name})" | |
| def postprocess(self, output: str): | |
| word_pattern = r"'?\w+(?:[ |.'-]\w+)*'?" | |
| triplet_pattern = f'({word_pattern}::{word_pattern}::{word_pattern})' | |
| processed = re.findall(triplet_pattern, re.sub(r'[<>]', '', output)) | |
| return processed | |