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

# Load the model and tokenizer
model_name = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Define the function to generate a response
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs.input_ids, 
        max_length=70,
        do_sample=True,
        temperature=0.6,
        top_p=0.9,
        repetition_penalty=1.2,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
    return response

# Persona-based response function
def persona_response(prompt, persona="You are a helpful talking dog that answers in short, simple phrases."):
    full_prompt = f"{persona}: {prompt}"
    return generate_response(full_prompt)

# Define Gradio interface function
def chat_interface(user_input, persona="You are a helpful talking dog that answers in short, simple phrases."):
    return persona_response(user_input, persona)

# Gradio interface setup
interface = gr.Interface(
    fn=chat_interface,
    inputs=["text", "text"],
    outputs="text",
    title="Simple Chatbot",
    description="Chat with the bot! Add a persona like 'I am a shopping assistant.'"
)

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