import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the DialoGPT-medium model and tokenizer (renamed to 1111)
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# Function for generating responses (renamed as 1111)
def chat_with_1111(input_text):
    # Encode the input text using the tokenizer
    input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')

    # Generate a response using the model (1111)
    chat_history_ids = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)

    # Decode and return the response
    response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return response

# Define the Gradio interface
iface = gr.Interface(
    fn=chat_with_1111,
    inputs=gr.Textbox(lines=2, placeholder="Ask 1111 a question..."),  # Input textbox
    outputs="text",  # Text output
    live=False,  # Do not send input live; only on button click
    title="Chat with 1111",
    description="Click the button to chat with 1111, an AI trained with DialoGPT medium!",
    examples=[["Hello!"]],  # Optional: example input for user to try out
    buttons=[gr.Button("Chat with 1111")]  # This adds a button
)

# Launch the interface (This will run on Aifaces when uploaded)
iface.launch()