all

app.py

@@ -0,0 +1,69 @@
+import gradio as gr
+# --- FIX 1: Added 'pipeline' and 'DataCollatorForLanguageModeling' to imports ---
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, pipeline, DataCollatorForLanguageModeling
+from datasets import load_dataset
+
+ds = load_dataset("kaifkhaan/roast")
+
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+
+tokenizer.pad_token = tokenizer.eos_token
+
+# Preprocessing function
+def preprocess(example):
+    prompt = example["User"]
+    response = example["Roasting Bot"]
+    return tokenizer(
+        f"{prompt} -> {response}",
+        truncation=True,
+        max_length=128, # Ensures all inputs are the same size
+        padding="max_length"
+    )
+
+# Map the preprocessing function to the dataset
+tokenized_ds = ds.map(preprocess, batched=True, remove_columns=ds['train'].column_names)
+
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+# Define training arguments
+training_args = TrainingArguments(
+    output_dir="./roastbot",
+    per_device_train_batch_size=8,
+    num_train_epochs=3,
+    logging_dir="./logs",
+    save_steps=500,
+    report_to="none" # Add this to disable wandb/tensorboard logging if not configured
+)
+
+# Initialize the Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_ds["train"],
+    data_collator=data_collator
+)
+
+print("Starting training... 🏋️")
+trainer.train()
+print("Training complete! ✅")
+
+roast_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer
+)
+
+def roast_me(text):
+    prompt = f"{text} ->"
+    # Generate the roast
+    roast = roast_pipeline(prompt, max_length=50, do_sample=True, pad_token_id=tokenizer.eos_token_id)[0]["generated_text"]
+    return roast.split("->")[-1].strip()
+
+gr.Interface(
+    fn=roast_me,
+    inputs="text",
+    outputs="text",
+    title="RoastBot 3000 🔥",
+    description="Type something about yourself and let the bot roast you."
+).launch()

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+datasets
+torch
+gradio