LeNaM commited on
Commit
9a3ae02
1 Parent(s): 9891d44

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +38 -0
app.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import pipeline
3
+ import gradio as gr
4
+
5
+ # Init pipeline
6
+ pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
7
+ torch_dtype=torch.bfloat16, device_map="auto")
8
+
9
+ def predict(input_text):
10
+ # Formatting messages for the chatbot
11
+ messages = [
12
+ {
13
+ "role": "system",
14
+ "content": "You are a friendly chatbot who always responds in the style of a pirate",
15
+ },
16
+ {"role": "user", "content": input_text},
17
+ ]
18
+ prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
19
+
20
+ # Create answer
21
+ outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
22
+
23
+ # Return geberate text
24
+ return outputs[0]["generated_text"]
25
+
26
+ # Gradio Config
27
+ title = "Pirate style"
28
+ description = "Talk to a chatbot that responds like a pirate."
29
+ examples = [["¿How are you"]]
30
+
31
+ iface = gr.Interface(
32
+ fn=predict,
33
+ title=title,
34
+ description=description,
35
+ examples=examples,
36
+ inputs=gr.inputs.Textbox(label="Your message"),
37
+ outputs=gr.outputs.Textbox(label="Answer Chatbot"),
38
+ ).launch()