import gradio as gr
from huggingface_hub import InferenceClient

# Define available models
MODELS = {
    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
    "GPT-2": "gpt2",
    "GPT-2 Medium": "gpt2-medium", 
    "DistilGPT-2": "distilgpt2",
    "German GPT-2": "german-nlp-group/german-gpt2",
    "German Wechsel GPT-2": "benjamin/gpt2-wechsel-german",
    "T5 Base": "t5-base",
    "T5 Large": "t5-large"
}

def create_inference_client(model_name):
    """Create an InferenceClient for the selected model."""
    try:
        return InferenceClient(model_name)
    except Exception as e:
        print(f"Error creating client for {model_name}: {e}")
        return None

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    model_name,
    max_tokens,
    temperature,
    top_p,
):
    """Generate response using selected model."""
    # Create client for selected model
    client = create_inference_client(MODELS[model_name])
    
    if not client:
        return "Error: Could not create model client."

    # Prepare chat history
    messages = [{"role": "system", "content": system_message}]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})

    # Generate response
    try:
        response = ""
        for message in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = message.choices[0].delta.content or ""
            response += token
            yield response
    except Exception as e:
        yield f"Error during generation: {e}"

def create_chat_interface():
    """Create Gradio ChatInterface with model selection."""
    demo = gr.ChatInterface(
        respond,
        additional_inputs=[
            gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
            gr.Dropdown(list(MODELS.keys()), value="Zephyr 7B Beta", label="Select Model"),
            gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
            gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p (nucleus sampling)",
            ),
        ]
    )
    return demo

if __name__ == "__main__":
    chat_interface = create_chat_interface()
    chat_interface.launch(share=True)