import gradio as gr
from huggingface_hub import InferenceClient

# Initialize client
client = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta")

def chat_function(message, history):
    messages = [
        {"role": "system", "content": "You are a helpful and professional AI assistant."}
    ]

    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})

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

    response = ""

    for message_obj in client.chat_completion(
        messages=messages,
        max_tokens=512,
        stream=True,
        temperature=0.7,
        top_p=0.95,
    ):
        token = message_obj.choices[0].delta.content or ""
        response += token
        yield response

demo = gr.ChatInterface(
    fn=chat_function,
    title="Premium AI Chatbot",
    description="A high-performance assistant hosted on Hugging Face.",
    examples=[
        "What are the benefits of RPA?",
        "How to optimize SIP investments?",
        "Explain financial planning."
    ],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch()