| 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", |
| "EmergentMethods/Phi-3-medium-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): |
| messages = [ |
| { |
| "role": "system", |
| "content": 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} |
| """) |
| }, |
| { |
| "role": "user", |
| "content": dedent(f"""\n |
| -------Text begin------- |
| {text} |
| -------Text end------- |
| """) |
| } |
| ] |
| return messages |
| |
| |
| def extract(self, text): |
| messages = self._get_messages(text) |
| pipe_output = self._generate(messages) |
| print("pipe_output", pipe_output) |
| print("pipe_output json", pipe_output[0]["generated_text"]) |
| return pipe_output[0]["generated_text"] |
|
|
