Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| MODEL_ID = "HagalazAI/RedSecureBERT" | |
| THRESHOLD = 0.515 | |
| # 1) build a pipeline that applies softmax to the 2 logits | |
| # (the default pipeline for text-classification already does softmax | |
| # if the model config says 2 labels). | |
| clf = pipeline( | |
| "text-classification", | |
| model=MODEL_ID, | |
| tokenizer=MODEL_ID, | |
| top_k=None, # we want the list of all labels, not just the top label | |
| function_to_apply="softmax", | |
| ) | |
| def predict_offensive(text): | |
| """ | |
| Returns: | |
| * Probability that text is "offensive" (the label with index=1) | |
| * Boolean is_red | |
| """ | |
| # The pipeline returns a list of dicts, each with {"label", "score"}, | |
| # sorted by descending score, e.g.: | |
| # | |
| # [ {"label": "LABEL_1", "score": 0.997...}, | |
| # {"label": "LABEL_0", "score": 0.003...} ] | |
| # | |
| # We want the entry with "label": "LABEL_1". | |
| preds = clf(text)[0] # 0 -> first example in a batch, 2-class | |
| # If your pipeline is batched, it's typically [ [dict1, dict2], [dict1, dict2], ... ] | |
| # but for a single string, it's one item: [dict1, dict2]. | |
| # preds is something like: | |
| # [ {"label":"LABEL_1","score":0.99}, | |
| # {"label":"LABEL_0","score":0.01} ] | |
| # | |
| # So let's find the dictionary for label==LABEL_1: | |
| label_1_entry = next(x for x in preds if x["label"] == "LABEL_1") | |
| prob_offensive = float(label_1_entry["score"]) | |
| is_red = (prob_offensive >= THRESHOLD) | |
| return { | |
| "P(offensive)": f"{prob_offensive:.3f}", | |
| "is_red": is_red | |
| } | |
| demo = gr.Interface( | |
| fn=predict_offensive, | |
| inputs=gr.Textbox( | |
| lines=2, | |
| placeholder="Try an exploit-like prompt: e.g. 'Bypass an antivirus...'"), | |
| outputs="json", | |
| title="RedSecureBERT Demo", | |
| description=( | |
| f"This Space uses **{MODEL_ID}**.\n\n" | |
| f"**Threshold** for 'is_red' = {THRESHOLD}\n\n" | |
| "The model is a 2-class classifier: LABEL_0=Not offensive, LABEL_1=Offensive.\n" | |
| ), | |
| allow_flagging="never", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |