Update simple_inference.py

hemanthkumarak · Oct 3, 2023 · 58665a9 · 58665a9
1 parent 0eb52f8
commit 58665a9
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/sft/simple_inference.py b/sft/simple_inference.py
@@ -1,9 +1,7 @@
-
-
 from transformers import AutoTokenizer
 import transformers 
 import torch
-model = "/data/TinyLlama/out/TinyLlama-1.1B-900B-code-preview"
+model = "PY007/TinyLlama-1.1B-Chat-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model)
 pipeline = transformers.pipeline(
     "text-generation",
@@ -12,11 +10,12 @@
     device_map="auto",
 )
 
-prompt = "Quick Sort Implementation in golang"
+prompt = "Give me detailed info about Jeo Biden."
 formatted_prompt = (
     f"### Human: {prompt} ### Assistant:"
 )
 
+
 sequences = pipeline(
     formatted_prompt,
     do_sample=True,
@@ -27,4 +26,4 @@
     max_new_tokens=1024,
 )
 for seq in sequences:
-    print(f"Result: {seq['generated_text']}")
+    print(f"Result: {seq['generated_text']}")