import os
import gradio as gr
from huggingface_hub import InferenceClient

MODEL_NAME = "google/gemma-2-9b-it"
HF_TOKEN = os.getenv("HF_TOKEN")

client = InferenceClient(
    model=MODEL_NAME,
    token=HF_TOKEN
)

def chat(message, history):
    messages = [{"role": "system", "content": "أنت مساعد ذكي."}]

    for u, a in history:
        if u:
            messages.append({"role": "user", "content": u})
        if a:
            messages.append({"role": "assistant", "content": a})

    messages.append({"role": "user", "content": message})

    response = client.chat_completion(
        messages=messages,
        max_tokens=300
    )

    return response.choices[0].message["content"]

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