import gradio as gr
from huggingface_hub import InferenceClient
import os

# Initialize the client with the token

client = InferenceClient("Dolly135/Pen_Model", token=os.getenv("token"))


# Modify the respond function to accept additional arguments
def respond(
    message: str,
    history: list,  # Simplified to accept Gradio's format
    system_message: str,
    temperature: float,
    top_p: float,
    max_new_tokens: int,  # New parameter for max tokens
):
    # Initialize messages with the system message
    messages = [{"role": "system", "content": system_message}]

    # Append history messages
    messages.extend(history)

    # Append the current user message
    messages.append({"role": "user", "content": message})

    try:
        response = ""
        for msg in client.text_generation(
            messages,
            stream=True,
            temperature=temperature,
            top_p=top_p,
            max_new_tokens=max_new_tokens,  # Pass max tokens to model
        ):
            token = msg.choices[0].delta.content
            response += token
            yield response
    except Exception as e:
        yield f"An error occurred: {str(e)}"

# Gradio setup using ChatInterface
if gr.__version__ >= '0.8.0':
    demo = gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Textbox(value="You are Pen.", label="System message"),
            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)"),
        ],
    )
else:
    demo = gr.ChatInterface(
        fn=respond,
        system_message=gr.Textbox(value="You are Pen.", label="System message"),
        max_tokens=gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        temperature=gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        top_p=gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    )

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