navyamehta commited on
Commit
38f9a42
·
verified ·
1 Parent(s): 296308b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -0
app.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
3
+
4
+ model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
5
+
6
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
7
+ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
8
+
9
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
10
+
11
+ def chat(message):
12
+ messages = [
13
+ {"role": "system", "content": "You are a helpful assistant."},
14
+ {"role": "user", "content": message}
15
+ ]
16
+ prompt = f"<|user|>\n{message}\n<|assistant|>\n"
17
+ output = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]["generated_text"]
18
+ return output.split("<|assistant|>\n")[-1].strip()
19
+
20
+ iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="TinyLlama Chat")
21
+ iface.launch()