Update README.md

README.md

@@ -1,15 +1,22 @@
 ```python
+
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-
-model_id = "nm-testing/TinyLlama-1.1B-Chat-v1.0-pruned50-24"
+prompt = "How to make banana bread?"
+formatted_prompt =  f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+model_id = "/home/mwiti/mwitiderrick/output_oneshot"
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-inputs = tokenizer("Hello my name is", return_tensors="pt")
-outputs = model.generate(**inputs, max_new_tokens=20)
+inputs = tokenizer(formatted_prompt, return_tensors="pt")
+outputs = model.generate(**inputs, max_new_tokens=200)
 print(tokenizer.batch_decode(outputs)[0])
+
 """
-<s> Hello my name is John. I am a student at the University of the University of the University of the University of the
+<s> <|im_start|>user
+How to make banana bread?<|im_end|>
+<|im_start|>assistant
+Banana bread is a delicious dessert that is made with bananas. Here is how to make banana bread:
 
+1. Firstly, you need to cut bananas into small pieces.
 """
 ```