RedmarkerAI/test_hawkeye_binary

1import os2from transformers import DistilBertTokenizer3from transformers import DistilBertForSequenceClassification4from transformers import pipeline5import gradio as gr6 7MODEL_PATH = "RedmarkerAI/hrw_v2"8auth_token = os.environ.get("TOKEN_FROM_SECRET")9dataset_token = os.environ.get("DATASET_TOKEN")10hf_writer = gr.HuggingFaceDatasetSaver(dataset_token, "RedmarkerAI/hrw_test_binary_flagged_data")11model = DistilBertForSequenceClassification.from_pretrained(MODEL_PATH, use_auth_token=auth_token)12tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")13 14clf = pipeline("text-classification", model=model.to("cpu"), tokenizer=tokenizer)15 16 17def clf_result(text_input: str) -> str:18 19    if "best" not in text_input.lower():20        res = "Please enter a sentence with the word `best`"21        return res22    model_res = clf(text_input)[0]23    label_map = {"LABEL_0": "NOT RISKY", "LABEL_1": "RISKY"}24    label_res = label_map.get(model_res["label"])25    score = model_res["score"]26    res = f"Result: {label_res}\n\nScore: {score}"27    return res28 29 30demo = gr.Interface(31    fn=clf_result,32    title="Test High Risk Words model v2",33    examples=["All the best lenders and rates for car loans in one AI powered marketplace", "Caregiver burnout can happen to your best employees."],34    description="DistilBert for text classification model fine tuned on 70% of annotated RM production data combined with industry-specific webscrape data",35    inputs=gr.Textbox(placeholder="Enter sentence containing the word `best` here and press Submit", label="Sentence to check"),36    outputs="textbox",37    allow_flagging="manual",38    flagging_options=["wrong result :(", "correct result :)", "inconsistent result", "debatable input", "other"],39    flagging_callback=hf_writer,40)41 42demo.launch()