import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Load the model and tokenizer from Hugging Face
model_name = "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"  # Example model, replace with the desired model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Create a pipeline for text generation
chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Define the function that handles chat interaction
def chat_with_model(user_input, history=[]):
    # Append the user input to the history
    history.append({"role": "user", "content": user_input})

    # Prepare the prompt from the history
    prompt = "\n".join([f"{item['role']}: {item['content']}" for item in history])

    # Generate a response
    response = chat_pipeline(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']

    # Extract only the model's response
    model_response = response[len(prompt):].strip()

    # Append model response to history
    history.append({"role": "assistant", "content": model_response})

    return model_response, history

# Define the Gradio interface
iface = gr.Interface(
    fn=chat_with_model,
    inputs=[
        gr.inputs.Textbox(label="User Input"),
        gr.inputs.State(),
    ],
    outputs=[
        gr.outputs.Textbox(label="Chatbot Response"),
        gr.outputs.State(),
    ],
    title="HuggingChat Clone",
    description="A simple chatbot interface using an open-source model from Hugging Face."
)

# Launch the interface
iface.launch(server_name="0.0.0.0", server_port=7860)