Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| from textwrap import dedent | |
| from huggingface_hub import InferenceClient | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| api_token = os.environ["HF_TOKEN"] | |
| endpoint_url = os.environ["HF_API_ENDPOINT"] | |
| # Initialize the client with your endpoint URL and token. | |
| client = InferenceClient( | |
| model=endpoint_url, | |
| token=api_token | |
| ) | |
| class Phi3InstructGraph: | |
| def __init__(self, model = "EmergentMethods/Phi-3-mini-4k-instruct-graph"): | |
| """ | |
| Initialize the Phi3InstructGraph with a specified model. | |
| """ | |
| self.model_path = model | |
| def _generate(self, messages): | |
| """ | |
| Generate a response from the model based on the provided messages. | |
| """ | |
| # Use the chat_completion method | |
| response = client.chat_completion( | |
| messages=messages, | |
| max_tokens=1024, | |
| ) | |
| # Access the generated text | |
| generated_text = response.choices[0].message.content | |
| return generated_text | |
| def _get_messages(self, text): | |
| """ | |
| Construct the message list for the chat model. | |
| """ | |
| context = dedent("""\n | |
| A chat between a curious user and an artificial intelligence Assistant. The Assistant is an expert at identifying entities and relationships in text. The Assistant responds in JSON output only. | |
| The User provides text in the format: | |
| -------Text begin------- | |
| <User provided text> | |
| -------Text end------- | |
| The Assistant follows the following steps before replying to the User: | |
| 1. **identify the most important entities** The Assistant identifies the most important entities in the text. These entities are listed in the JSON output under the key "nodes", they follow the structure of a list of dictionaries where each dict is: | |
| "nodes":[{"id": <entity N>, "type": <type>, "detailed_type": <detailed type>}, ...] | |
| where "type": <type> is a broad categorization of the entity. "detailed type": <detailed_type> is a very descriptive categorization of the entity. | |
| 2. **determine relationships** The Assistant uses the text between -------Text begin------- and -------Text end------- to determine the relationships between the entities identified in the "nodes" list defined above. These relationships are called "edges" and they follow the structure of: | |
| "edges":[{"from": <entity 1>, "to": <entity 2>, "label": <relationship>}, ...] | |
| The <entity N> must correspond to the "id" of an entity in the "nodes" list. | |
| The Assistant never repeats the same node twice. The Assistant never repeats the same edge twice. | |
| The Assistant responds to the User in JSON only, according to the following JSON schema: | |
| {"type":"object","properties":{"nodes":{"type":"array","items":{"type":"object","properties":{"id":{"type":"string"},"type":{"type":"string"},"detailed_type":{"type":"string"}},"required":["id","type","detailed_type"],"additionalProperties":false}},"edges":{"type":"array","items":{"type":"object","properties":{"from":{"type":"string"},"to":{"type":"string"},"label":{"type":"string"}},"required":["from","to","label"],"additionalProperties":false}}},"required":["nodes","edges"],"additionalProperties":false} | |
| """) | |
| user_message = dedent(f"""\n | |
| -------Text begin------- | |
| {text} | |
| -------Text end------- | |
| """) | |
| # if self.model_path == "EmergentMethods/Phi-3-medium-128k-instruct-graph": | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": context | |
| }, | |
| { | |
| "role": "user", | |
| "content": user_message | |
| } | |
| ] | |
| # else: | |
| # # TODO: update for other models | |
| return messages | |
| def extract(self, text): | |
| """ | |
| Extract knowledge graph from text | |
| """ | |
| messages = self._get_messages(text) | |
| generated_text = self._generate(messages) | |
| return generated_text | |