Update README.md

README.md

@@ -55,11 +55,47 @@ Use the code in [this repository](https://github.com/chtmp223/suri) for training
 | optim                            | adamw_torch  |
 | per_device_train_batch_size      | 1            |
 
-
-#### 🤗 Software
+#### Software
 
 Training code is adapted from [Alignment Handbook](https://github.com/huggingface/alignment-handbook) and [Trl](https://github.com/huggingface/trl).
 
+## 🤗 Inference
+
+```
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
+from datasets import load_dataset
+import torch
+os.environ["TOKENIZERS_PARALLELISM"] = "False"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.cuda.empty_cache()
+
+model_name = "chtmp223/suri-i-orpo"
+base_model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+config = PeftConfig.from_pretrained(model_name)
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name).to(device)
+model = PeftModel.from_pretrained(base_model, model_name).to(device)
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+prompt = [
+  {
+      "role": "user",
+      "content": user_prompt, 
+  }
+]
+input_context = tokenizer.apply_chat_template(
+  prompt, add_generation_prompt=True, tokenize=False
+)
+input_ids = tokenizer.encode(
+  input_context, return_tensors="pt", add_special_tokens=False
+).to(model.device)
+output = model.generate(
+  input_ids, max_length=10000, do_sample=True, use_cache=True
+).cpu()
+
+print(tokenizer.decode(output[0]))
+```
+
+
 ## 📜 Citation 
 
 ```