import gradio as gr
from transformers import pipeline

# We switched to a 'community' version that doesn't have a locked door
pipe = pipeline("text-generation", model="unsloth/gemma-2-2b-it-bnb-4bit")

def chat_with_ai(prompt):
    messages = [{"role": "user", "content": prompt}]
    output = pipe(messages, max_new_tokens=150)
    return output[0]["generated_text"][-1]["content"]

demo = gr.Interface(
    fn=chat_with_ai, 
    inputs=gr.Textbox(label="Ask Gemma anything!"), 
    outputs="text",
    title="My Gemma Bot"
)

demo.launch()