import os
from transformers import DistilBertTokenizer
from transformers import DistilBertForSequenceClassification
from transformers import pipeline
import gradio as gr

MODEL_PATH = "RedmarkerAI/hrw_v2"
auth_token = os.environ.get("TOKEN_FROM_SECRET")
dataset_token = os.environ.get("DATASET_TOKEN")
hf_writer = gr.HuggingFaceDatasetSaver(dataset_token, "RedmarkerAI/hrw_test_binary_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_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 model v2",
    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()