Update app.py

app.py

@@ -1,74 +1,78 @@
-# import torch
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-# print(f"Using device: {device}")
 
+import torch
+import pandas as pd
+import re
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+# Set device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Load model
 model_name = "rajistics/informal_formal_style_transfer"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-import pandas as pd
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
 
+# Load dataset
 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()
 
+# Random samples from dataset
 gradio_examples = [["[Gen-Z Example] " + row["slang"]] for _, row in df.sample(3, random_state=1).iterrows()]
-gradio_examples
 
+# Manual 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"]
+    ["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
+# Clean output
 def clean_output(output: str):
-    return re.sub(r"(?i)make this sentence more formal", "", output).strip()
+    return re.sub(r"(?i)make this sentence more formal", "", output).strip()
 
-# Main function
+# Model inference
 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)
+    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)
 
+# Gradio app
 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"
+    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!"
+    ),
+    article=(
+        "**Project by Jonathan Friedman**  \n"
+        "📌 Task: Text Generation (Formality Transfer using Sequence-to-Sequence)\n"
+        "🧠 Model: rajistics/informal_formal_style_transfer\n"
+        "📚 Dataset: thesherrycode/gen-z-slangs-translation\n"
+        "🛠️ Tech Stack: Hugging Face Transformers, Gradio"
+    ),
+    examples=examples + gradio_examples,
+    theme="soft"
 )
 
-
-iface.launch()
+iface.launch()