# import torch
# device = "cuda" if torch.cuda.is_available() else "cpu"
# print(f"Using device: {device}")

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_name = "rajistics/informal_formal_style_transfer"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
import pandas as pd

df = pd.read_csv("hf://datasets/thesherrycode/gen-z-slangs-translation/gen_z_slangs_translation.csv")

# Rename the actual columns
df = df[["Gen-Z Slang", "Plain English"]].dropna().drop_duplicates()
df.columns = ["slang", "formal"]
df.head()

gradio_examples = [["[Gen-Z Example] " + row["slang"]] for _, row in df.sample(3, random_state=1).iterrows()]
gradio_examples

examples = [
    ["hey, can u send me the stuff by tonight?"],
    ["yo sorry i missed the call, was busy"],
    ["lemme know if ur free tmrw to chat abt the thing"],
    ["bro the file’s messed up, fix it asap pls"],
    ["i'm out rn, text u later"]
]

# Clean the model output
def clean_output(output: str):
    return re.sub(r"(?i)make this sentence more formal", "", output).strip()

# Main function
def make_formal(text):
    if not text.strip():
        return "⚠️ Please enter some text."
    prompt = "[Casual] " + text.strip() + " [Formal]"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(device)
    outputs = model.generate(**inputs, max_new_tokens=100)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return clean_output(result)

iface = gr.Interface(
    fn=make_formal,
    inputs=gr.Textbox(
        lines=3,
        label="🗣️ Your Slang / Casual Text",
        placeholder="e.g., yo can u help me out real quick?"
    ),
    outputs=gr.Textbox(
        label="📄 Formal & Polished Version",
        lines=4,
        interactive=True
    ),
    title="💬 Text Polisher: From Slang to Formal",
    description=(
        "Transform casual, Gen-Z slang, or unpolished English into clear, professional language. 🧠✨\n\n"
        "This demo uses a text generation model to rewrite input sentences with improved formality — great for school, work, or writing more professionally.\n\n"
        "✍️ The output is editable — feel free to tweak before using/copying!\n\n"
    ),
    article=(
        "**Project by Jonathan Friedman**  \n"
        "📌 **Task:** Text Generation (Formality Transfer using Sequence-to-Sequence)  \n"
        "🧠 **Model:** rajistics/informal_formal_style_transfer  \n"
        "📚 **Dataset Used:** thesherrycode/gen-z-slangs-translation (used to generate real-world Gen-Z slang examples)  \n"
        "🛠️ **Tech Stack:** 🤗 Transformers, Hugging Face Datasets, Gradio  \n"
    ),
    examples=examples + gradio_examples,
    theme="soft"
)


iface.launch()