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 conversational text generation chatbot.", | |
}, | |
{ | |
"role": "user", | |
"content": "Hello! I am a chatbot designed to generate conversational text. How can I assist you today?", | |
} | |
] | |
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 = "Conversation style" | |
description = "Talk to a chatbot that responds like a conversational chat." | |
examples = [["¿How are you"]] | |
iface = gr.Interface( | |
fn=predict, | |
title=title, | |
description=description, | |
examples=examples, | |
inputs=gr.Textbox(label="Your message"), | |
outputs=gr.Textbox(label="Answer Chatbot"), | |
).launch() | |