my-gemma-bot / app.py
abhy60098's picture
Update app.py
5c1cc75 verified
Raw
History Blame Contribute Delete
541 Bytes
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()