Initial application file

app.py

@@ -0,0 +1,25 @@
+from transformers import AutoTokenizer
+from transformers import AutoModelForSequenceClassification
+from transformers import DataCollatorWithPadding
+from transformers import TrainingArguments, Trainer
+import gradio as gr
+
+tokenizer = AutoTokenizer.from_pretrained("smallbenchnlp/roberta-small")
+training_args = TrainingArguments(output_dir="roberta-small-pun-detector", evaluation_strategy="epoch")
+model = AutoModelForSequenceClassification.from_pretrained("frostymelonade/roberta-small-pun-detector", num_labels=2)
+data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    data_collator=data_collator,
+    tokenizer=tokenizer,
+)
+
+def classify_pun(text):
+    inputs = [tokenizer(text, truncation=True)]
+    predictions = trainer.predict(inputs)
+    label = "Pun" if predictions[0][0][0] < predictions[0][0][1] else "Not a pun"
+    return label
+
+gr.Interface(fn=classify_pun, inputs=["text"], outputs=["text"]).launch()