Views
No views yet
1from peft import PeftModel, PeftConfig
2from transformers import AutoModelForCausalLM, AutoTokenizer
3import torch
4from transformers import StoppingCriteria
5
6config = PeftConfig.from_pretrained("Mit1208/phi-2-universal-NER")
7base_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2",device_map="auto", trust_remote_code=True)
8model = PeftModel.from_pretrained(base_model, "Mit1208/phi-2-universal-NER", trust_remote_code=True)
9tokenizer = AutoTokenizer.from_pretrained("Mit1208/phi-2-universal-NER", trust_remote_code=True)
10
11device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
13conversations = [ { "from": "human", "value": "Text: Mit Patel here from India"}, {"from": "gpt", "value": "I've read this text."},
14 {"from":"human", "value":"what is a name of the person in the text?"}]
15inference_text = tokenizer.apply_chat_template(conversations, tokenize=False) + '<|im_start|>gpt:\n'
16inputs = tokenizer(inference_text, return_tensors="pt", return_attention_mask=False).to(device)
17
18class EosListStoppingCriteria(StoppingCriteria):
19 def __init__(self, eos_sequence = tokenizer.encode("<|im_end|>")):
20 self.eos_sequence = eos_sequence
21
22 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
23 last_ids = input_ids[:,-len(self.eos_sequence):].tolist()
24 return self.eos_sequence in last_ids
25
26outputs = model.generate(**inputs, max_length=512, pad_token_id= tokenizer.eos_token_id,
27 stopping_criteria = [EosListStoppingCriteria()])
28
29text = tokenizer.batch_decode(outputs)[0]
30
31print(text)
32
33# Output
34'''
35<|im_start|>human
36Text: Mit Patel here from India<|im_end|>
37<|im_start|>gpt
38I've read this text.<|im_end|>
39<|im_start|>human
40what is a name of the person in the text?<|im_end|>
41<|im_start|>gpt:
42["Mit Patel"]<|im_end|>
43'''