import os
from huggingface_hub import InferenceClient
import gradio as gr

client = InferenceClient(api_key=os.environ["HF_API_KEY"])

def chat(message, history):
    r = client.chat_completion(
        model="google/gemma-2-2b-it",
        messages=[
            {"role": "user", "content": message}
        ],
        max_tokens=300
    )
    return r.choices[0].message["content"]

gr.ChatInterface(fn=chat).launch()