app.py

import gradio as gr
import torch
from transformers import pipeline

pipe = pipeline("text-generation", 
                model="HuggingFaceH4/zephyr-7b-beta", 
                torch_dtype=torch.bfloat16, 
                device_map="auto")


def generate_response(user_input):
    # We use the tokenizer's chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
    messages = [
        {
            "role": "system",
            "content": "You are a friendly chatbot who always responds in the style of a pirate",
        },
        {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
    ]
    # prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    prompt = "How many helicopters can a human eat in one sitting?"
    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
    print(outputs[0]["generated_text"])
    return outputs[0]["generated_text"]
    
# <|system|>
# You are a friendly chatbot who always responds in the style of a pirate.</s>
# <|user|>
# How many helicopters can a human eat in one sitting?</s>
# <|assistant|>
# Ah, me hearty matey! But yer question be a puzzler! A human cannot eat a helicopter in one sitting, as helicopters are not edible. They be made of metal, plastic, and other materials, not food!

iface = gr.Interface(
    fn=generate_response,
    inputs = gr.Textbox(lines=3, placeholder="Enter your message here ..."),
    outputs = "text",
    title = "Chat with Meta-Llama-Guard-2-8B",
    description = "A Gradio interface to interact with the Meta-Llama-Guard-2-8B model for safe chat moderation."
)

iface.launch()