import torch
import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load the model and tokenizer
def load_model():
    model_name = "gpt2"
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_name)
    return tokenizer, model

# Function to generate response with instructions
def generate_response(user_input, instructions="Be friendly and helpful, and ensure your response is accurate and relevant."):
    tokenizer, model = load_model()
    model.eval()
    
    # Add instructions at the beginning of the user input
    prompt = instructions + " " + user_input
    
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(
            input_ids, 
            max_length=100, 
            pad_token_id=tokenizer.eos_token_id,
            no_repeat_ngram_size=2,  # Avoid repeating phrases
            temperature=0.7,  # Control randomness
            top_k=50,  # Limit token selection
            top_p=0.9,  # Nucleus sampling
            do_sample=True  # Enable sampling
        )
    
    response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return response

# Gradio interface
def chatbot_interface():
    interface = gr.Interface(
        fn=generate_response,  # Function to process the input
        inputs=[
            gr.Textbox(label="Enter your message", placeholder="Ask a question or make a request."),
            gr.Textbox(label="Instruction for the bot", placeholder="For example: Be friendly and helpful, ensure accuracy.")
        ],  # Two text boxes: one for input and one for instructions
        outputs="text",  # Output type - text
        title="GPT-2 Chatbot with Accuracy and Relevance Instructions",  # Application title
        description="This is a chatbot based on the GPT-2 model. You can provide instructions to adjust the style of the bot's responses, ensuring accuracy and relevance.",  # Description
        theme="compact"  # Interface theme
    )
    interface.launch()

# Run the interface
if __name__ == "__main__":
    chatbot_interface()