app.py

"""
# Inference

import gradio as gr

app = gr.load(
    "google/gemma-2-2b-it",
    src = "models",
    inputs = [gr.Textbox(label = "Input")],
    outputs = [gr.Textbox(label = "Output")],
    title = "Google Gemma",
    description = "Inference",
    examples = [
        ["Hello, World."]
    ]
).launch()
"""
"""
# Pipeline

import gradio as gr
from transformers import pipeline

pipe = pipeline(model = "google/gemma-2-2b-it")

def fn(input):
    output = pipe(
        input,
        max_new_tokens = 2048
    )
    return output[0]["generated_text"]#[len(input):]

app = gr.Interface(
    fn = fn,
    inputs = [gr.Textbox(label = "Input")],
    outputs = [gr.Textbox(label = "Output")],
    title = "Google Gemma",
    description = "Pipeline",
    examples = [
        ["Hello, World."]
    ]
).launch()
"""

import gradio as gr
from huggingface_hub import InferenceClient
import os

hf_token = os.getenv("HF_TOKEN")
client = InferenceClient(api_key=hf_token)

def chatbot(input_text, history):
    # Prepare the conversation messages
    messages = [{"role": "user", "content": input_text}]

    # Add conversation history (if exists)
    if history:
        for user_input, bot_response in history:
            messages.append({"role": "user", "content": user_input})
            messages.append({"role": "assistant", "content": bot_response})

    # Generate model response
    stream = client.chat.completions.create(
        model="google/gemma-2-2b-it", 
        messages=messages, 
        temperature=0.5,
        max_tokens=2048,
        top_p=0.7,
        stream=True
    )

    # Collect response from the model
    bot_response = ""
    for chunk in stream:
        bot_response += chunk.choices[0].delta.content

    # Update the conversation history
    history.append((input_text, bot_response))
    return bot_response, history

# Create Gradio Interface
with gr.Blocks() as demo:
    chatbot_ui = gr.Chatbot(label="Gemma Chatbot").style(height=400)
    text_input = gr.Textbox(show_label=False, placeholder="Type your message here...")
    state = gr.State([])  # Keeps track of conversation history

    def user_input_handler(user_input, chatbot_ui, state):
        bot_response, state = chatbot(user_input, state)
        chatbot_ui.append((user_input, bot_response))
        return chatbot_ui, state

    text_input.submit(user_input_handler, [text_input, chatbot_ui, state], [chatbot_ui, state])

# Run the application
demo.launch()