CyberGuard / app.py
RohanSardar's picture
Update app.py
40ffef6 verified
raw
history blame contribute delete
894 Bytes
from transformers import BertForSequenceClassification, BertTokenizerFast, pipeline
import gradio as gr
model_path = "indiaai-text-classification-model"
model = BertForSequenceClassification.from_pretrained(model_path)
tokenizer = BertTokenizerFast.from_pretrained(model_path)
nlp = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
def classify_text(input_text):
result = nlp(input_text)
label = result[0]['label']
score = result[0]['score']
output = f"**Prediction:** {label}\n\n**Confidence Score:** {score:.5f}"
return output
interface = gr.Interface(
fn=classify_text,
inputs=gr.Textbox(lines=2, placeholder="Enter your complaint", label="Input"),
outputs=gr.Markdown(),
title="INDIAai CyberGuard",
description="Categorizes cyber complaints based on the victim, type of fraud, and other relevant parameters.",
)
interface.launch()