Update README.md

README.md

@@ -44,8 +44,51 @@ The following hyperparameters were used during training:
 - lr_scheduler_type: cosine
 - training_steps: 1000
 
-### Training results
-
+### Inference Code
+
+```python
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from transformers import StoppingCriteria
+
+config = PeftConfig.from_pretrained("Mit1208/phi-2-universal-NER")
+base_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2",device_map="auto", trust_remote_code=True)
+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)
+
+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)
+
+class EosListStoppingCriteria(StoppingCriteria):
+    def __init__(self, eos_sequence = tokenizer.encode("<|im_end|>")):
+        self.eos_sequence = eos_sequence
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        last_ids = input_ids[:,-len(self.eos_sequence):].tolist()
+        return self.eos_sequence in last_ids
+
+outputs = model.generate(**inputs, max_length=512, pad_token_id= tokenizer.eos_token_id,
+            stopping_criteria = [EosListStoppingCriteria()])
+
+text = tokenizer.batch_decode(outputs)[0]
+
+print(text)
+
+# Output
+'''
+<|im_start|>human
+Text: Mit Patel here from India<|im_end|>
+<|im_start|>gpt
+I've read this text.<|im_end|>
+<|im_start|>human
+what is a name of the person in the text?<|im_end|>
+<|im_start|>gpt:
+["Mit Patel"]<|im_end|>
+'''
+```
 
 
 ### Framework versions