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

# === Dossier contenant le modèle fine-tuné ===
MODEL_DIR = "model"

# === Chargement du modèle GPT-Neo ===
print("🔁 Loading GPT-Neo model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForCausalLM.from_pretrained(MODEL_DIR)

# === Ajouter un token spécial de stop ===
stop_token = "<|stop|>"
tokenizer.add_special_tokens({'additional_special_tokens': [stop_token]})
model.resize_token_embeddings(len(tokenizer))

# === Création de la pipeline ===
device = 0 if torch.cuda.is_available() else -1
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device
)

# === Fonction de réponse sarcastique ===
def respond(prompt):
    full_input = f"### Prompt: {prompt}\n### Response (sarcastic):"
    output = generator(
        full_input,
        max_new_tokens=60,
        do_sample=True,
        temperature=0.9,
        top_k=50,
        top_p=0.95,
        repetition_penalty=1.2,
        eos_token_id=tokenizer.convert_tokens_to_ids(stop_token),
    )
    raw_text = output[0]["generated_text"]
    response = raw_text.replace(full_input, "")
    response = re.split(re.escape(stop_token), response)[0].strip()
    if response.endswith(("and", "or", "but", "because")):
        response = response.rsplit(" ", 1)[0] + "..."
    return response

# === Interface Gradio ===
demo = gr.Interface(
    fn=respond,
    inputs=gr.Textbox(lines=2, placeholder="Ask something ridiculous..."),
    outputs="text",
    title="S.N.A.R.K – Sarcastic Neural Algorithm",
    description="A fine-tuned GPT-Neo model that delivers sharp and witty sarcasm on demand. Proceed at your own risk."
)

demo.launch()