Update code to generate text using TextStreamer

README.md

@@ -100,9 +100,15 @@ top_k: 5
 ## Usage Example
 
 ```python
-from transformers import pipeline
+from transformers import pipeline, TextStreamer, AutoModelForCausalLM, AutoTokenizer
+import torch
 
-generate = pipeline("text-generation", "Felladrin/TinyMistral-248M-Chat-v3")
+model_path = "Felladrin/TinyMistral-248M-Chat-v3"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
+streamer = TextStreamer(tokenizer)
+generate = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
 messages = [
     {
@@ -123,16 +129,27 @@ messages = [
     },
 ]
 
-prompt = generate.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-output = generate(
-    prompt,
-    max_new_tokens=256,
-    penalty_alpha=0.5,
-    top_k=5,
+prompt = tokenizer.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
 )
 
-print(output[0]["generated_text"])
+inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+outputs = model.generate(
+    inputs.input_ids,
+    attention_mask=inputs.attention_mask,
+    max_length=tokenizer.model_max_length,
+    streamer=streamer,
+    eos_token_id=tokenizer.eos_token_id,
+    pad_token_id=tokenizer.pad_token_id,
+    do_sample=True,
+    temperature=0.6,
+    top_p=0.8,
+    top_k=0,
+    min_p=0.1,
+    typical_p=0.2,
+    repetition_penalty=1.176,
+)
 ```
 
 ## How it was trained