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()