Spaces:
Running
Running
| import torch | |
| import sentencepiece | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline | |
| from langchain import PromptTemplate, LLMChain, HuggingFacePipeline | |
| import ast | |
| class Models(): | |
| def __init__(self) -> None: | |
| self.template = """ | |
| A virtual assistant answers questions from a user based on the provided text. | |
| USER: Text: {input_text} | |
| ASSISTANT: I’ve read this text. | |
| USER: What describes {entity_type} in the text? | |
| ASSISTANT: | |
| """ | |
| self.load_trained_models() | |
| def load_trained_models(self): | |
| #is it best to keep in memory why not pickle? | |
| checkpoint = "Universal-NER/UniNER-7B-all" | |
| ner_model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float16, offload_folder="offload", offload_state_dict = True) | |
| tokenizer = AutoTokenizer.from_pretrained("Universal-NER/UniNER-7B-all", use_fast=False, padding="max_length") | |
| pipeline = pipeline( | |
| "text-generation", #task | |
| model=ner_model, | |
| max_length=1000, | |
| tokenizer=tokenizer, | |
| trust_remote_code=True, | |
| do_sample=True, | |
| top_k=10, | |
| num_return_sequences=1 | |
| ) | |
| self.llm = HuggingFacePipeline(pipeline = pipeline, model_kwargs = {'temperature':0}) | |
| self.prompt = PromptTemplate(template=self.template, input_variables=["input_text","entity_type"]) | |
| self.llm_chain = LLMChain(prompt=self.prompt, llm=self.llm) | |
| def extract_ner(self, context, entity_type): | |
| return ast.literal_eval(self.llm_chain.run({"input_text":context,"entity_type":entity_type})) | |
| def get_ner(self, clean_lines, entity): | |
| tokens = [] | |
| try_num = 0 | |
| while try_num < 5 and tokens == []: | |
| tokens = self.extract_ner(' '.join(clean_lines), entity) | |
| if len(tokens) == 0: | |
| raise ValueError("Couldnt extract {entity}") | |
| return tokens | |