Opsifiz's picture
Model V0.2.1
a362f30
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import contractions
import html
import re
# Load model and tokenizer
model_path = "my_model" # directory with your trained model
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
# Labels
id2label = {
0: 'Anxiety',
1: 'BPD',
2: 'Normal',
3: 'Stress',
4: 'Suicidal',
5: 'Bipolar',
6: 'Depression',
7: 'Mentalillness',
8: 'Schizophrenia'
}
def clean_text(text):
text = str(text)
text = text.lower()
text = contractions.fix(text)
text = html.unescape(text)
text = re.sub(r'http\S+', '', text) # Remove URLs
text = re.sub(r'[^a-zA-Z\s]', '', text) # Remove special characters
text = re.sub(r'\s+', ' ', text).strip()
return text
# Inference function
def predict(text):
text = clean_text(text)
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
pred_id = outputs.logits.argmax(dim=1).item()
return id2label[pred_id]
# Gradio UI
demo = gr.Interface(
fn=predict,
inputs=gr.Textbox(lines=3, placeholder="Enter text here..."),
outputs=gr.Label(),
title="Text - Emotion Classifier",
description="Predicts emotion category based on text input."
)
if __name__ == "__main__":
demo.launch()