| | import torch |
| | from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
| | from textwrap import dedent |
| | from huggingface_hub import login |
| | import os |
| | from dotenv import load_dotenv |
| |
|
| | load_dotenv() |
| | login( |
| | token=os.environ["HF_TOKEN"], |
| | ) |
| |
|
| | MODEL_LIST = [ |
| | "EmergentMethods/Phi-3-mini-4k-instruct-graph", |
| | "EmergentMethods/Phi-3-mini-128k-instruct-graph", |
| | |
| | ] |
| |
|
| | torch.random.manual_seed(0) |
| |
|
| | class Phi3InstructGraph: |
| | def __init__(self, model = "EmergentMethods/Phi-3-mini-4k-instruct-graph"): |
| | if model not in MODEL_LIST: |
| | raise ValueError(f"model must be one of {MODEL_LIST}") |
| | |
| | self.model_path = model |
| | self.model = AutoModelForCausalLM.from_pretrained( |
| | self.model_path, |
| | device_map="cuda", |
| | torch_dtype="auto", |
| | trust_remote_code=True, |
| | ) |
| | self.tokenizer = AutoTokenizer.from_pretrained(self.model_path) |
| | self.pipe = pipeline( |
| | "text-generation", |
| | model=self.model, |
| | tokenizer=self.tokenizer, |
| | ) |
| |
|
| | def _generate(self, messages): |
| | generation_args = { |
| | "max_new_tokens": 2000, |
| | "return_full_text": False, |
| | "temperature": 0.1, |
| | "do_sample": False, |
| | } |
| |
|
| | return self.pipe(messages, **generation_args) |
| |
|
| | def _get_messages(self, text): |
| | 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": "user", |
| | "content": f"{context}\n Input: {user_message}", |
| | } |
| | ] |
| | return messages |
| | else: |
| | messages = [ |
| | { |
| | "role": "system", |
| | "content": context |
| | }, |
| | { |
| | "role": "user", |
| | "content": user_message |
| | } |
| | ] |
| | return messages |
| | |
| | |
| | def extract(self, text): |
| | messages = self._get_messages(text) |
| | pipe_output = self._generate(messages) |
| | |
| | return pipe_output[0]["generated_text"] |
| |
|
