Create app.py

app.py

@@ -0,0 +1,47 @@
+import os
+from transformers import DistilBertTokenizer
+from transformers import DistilBertForSequenceClassification
+from transformers import pipeline
+import gradio as gr
+
+MODEL_PATH = "RedmarkerAI/hrw_multi_generous_v1"
+auth_token = os.environ.get("TOKEN_FROM_SECRET")
+dataset_token = os.environ.get("DATASET_TOKEN")
+hf_writer = gr.HuggingFaceDatasetSaver(dataset_token, "hrw_test_multiclass_flagged_data")
+model = DistilBertForSequenceClassification.from_pretrained(MODEL_PATH, use_auth_token=auth_token)
+tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+
+clf = pipeline("text-classification", model=model.to("cpu"), tokenizer=tokenizer)
+
+
+def clf_result(text_input: str) -> str:
+
+    if "best" not in text_input.lower():
+        res = "Please enter a sentence with the word `best`"
+        return res
+    model_res = clf(text_input)[0]
+    # label_map = {"LABEL_0": "NOT RISKY", "LABEL_1": "RISKY"}
+    label_map = {
+    "0": "Not Risky",
+    "1": "Possibly Risky",
+    "2": "Risky"
+  }
+    label_res = label_map.get(model_res["label"])
+    score = model_res["score"]
+    res = f"Result: {label_res}\n\nScore: {score}"
+    return res
+
+
+demo = gr.Interface(
+    fn=clf_result,
+    title="Test High Risk Words Multiclass model",
+    examples=["All the best lenders and rates for car loans in one AI powered marketplace", "Caregiver burnout can happen to your best employees."],
+    description="DistilBert for text classification model fine tuned on 70% of annotated RM production data combined with industry-specific webscrape data",
+    inputs=gr.Textbox(placeholder="Enter sentence containing the word `best` here and press Submit", label="Sentence to check"),
+    outputs="textbox",
+    allow_flagging="manual",
+    flagging_options=["wrong result :(", "correct result :)", "inconsistent result", "debatable input", "other"],
+    flagging_callback=hf_writer,
+)
+
+demo.launch()