Update README.md

README.md

@@ -56,10 +56,12 @@ base_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2",device_map="
 model = PeftModel.from_pretrained(base_model, "Mit1208/phi-2-universal-NER", trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained("Mit1208/phi-2-universal-NER", trust_remote_code=True)
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 conversations = [ { "from": "human", "value": "Text: Mit Patel here from India"}, {"from": "gpt", "value": "I've read this text."}, 
                    {"from":"human", "value":"what is a name of the person in the text?"}]
 inference_text = tokenizer.apply_chat_template(conversations, tokenize=False) + '<|im_start|>gpt:\n'
-inputs = tokenizer(inference_text, return_tensors="pt", return_attention_mask=False)
+inputs = tokenizer(inference_text, return_tensors="pt", return_attention_mask=False).to(device)
 
 class EosListStoppingCriteria(StoppingCriteria):
     def __init__(self, eos_sequence = tokenizer.encode("<|im_end|>")):