Spaces:
Sleeping
Sleeping
import torch | |
from transformers import pipeline | |
import gradio as gr | |
# Init pipeline | |
pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", | |
torch_dtype=torch.bfloat16, device_map="auto") | |
def predict(input_text): | |
# Formatting messages for the chatbot | |
messages = [ | |
{ | |
"role": "system", | |
"content": "You are a friendly chatbot who always responds in the style of a pirate", | |
}, | |
{"role": "user", "content": input_text}, | |
] | |
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
# Create answer | |
outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95) | |
# Return geberate text | |
return outputs[0]["generated_text"] | |
# Gradio Config | |
title = "Pirate style" | |
description = "Talk to a chatbot that responds like a pirate." | |
examples = [["¿How are you"]] | |
iface = gr.Interface( | |
fn=predict, | |
title=title, | |
description=description, | |
examples=examples, | |
inputs=gr.inputs.Textbox(label="Your message"), | |
outputs=gr.outputs.Textbox(label="Answer Chatbot"), | |
).launch() | |