import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the Hugging Face Inference API client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Function to generate responses using the Hugging Face model
def respond(message, history):
    # Prepare the conversation history
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    # Generate a response using the model
    response = ""
    for chunk in client.chat_completion(
        messages,
        max_tokens=512,  # Adjust as needed
        stream=True,
        temperature=0.7,  # Adjust for creativity
        top_p=0.95,       # Adjust for diversity
    ):
        token = chunk.choices[0].delta.content
        response += token

    return response

# Gradio ChatInterface
demo = gr.ChatInterface(
    respond,
    title="Simple AI Chatbot",
    description="Ask me anything!",
    examples=[
        "What is the capital of France?",
        "Explain quantum computing in simple terms.",
        "Write a short poem about the ocean."
    ]
)

# Launch the app
if __name__ == "__main__":
    demo.launch()