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

# ✅ 1. Load GPT-Neo model
model_name = "EleutherAI/gpt-neo-1.3B"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)

# ✅ 2. Chatbot-style Text Generator
def chat_with_gptneo(history, user_input):
    # Combine chat history for context (optional, keep recent exchanges)
    if history:
        past = "\n".join([f"User: {msg[0]}\nBot: {msg[1]}" for msg in history])
        prompt = f"{past}\nUser: {user_input}\nBot:"
    else:
        prompt = f"User: {user_input}\nBot:"

    inputs = tokenizer(prompt, return_tensors='pt').to(device)

    outputs = model.generate(
        inputs['input_ids'],
        max_new_tokens=200,
        do_sample=True,
        temperature=0.85,
        top_k=50,
        top_p=0.95,
        repetition_penalty=1.2,
        no_repeat_ngram_size=3,
        pad_token_id=tokenizer.eos_token_id,
        early_stopping=True
    )

    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = full_output.replace(prompt, "").strip()
    return response

# ✅ 3. Gradio Chat Interface
with gr.Blocks(title="GPT-Neo Chatbot") as demo:
    gr.Markdown("## 🤖 GPT-Neo 1.3B Chatbot\nAsk me anything!")
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="Type your prompt here...", label="Your Prompt")
    clear = gr.Button("Clear")

    history = gr.State([])

    def respond(user_message, chat_history):
        bot_message = chat_with_gptneo(chat_history, user_message)
        chat_history.append((user_message, bot_message))
        return "", chat_history

    msg.submit(respond, [msg, history], [msg, chatbot])
    clear.click(lambda: ([], []), None, [history, chatbot])

# ✅ 4. Launch the app
demo.launch()