Update app.py

app.py

@@ -2,10 +2,13 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Load Pretrained Model & Tokenizer (Ensure this is a fine-tuned model)
-MODEL_NAME = "xlm-roberta-base"
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=5)  # Adjust num_labels as per training
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+# Load Fine-Tuned Model & Tokenizer (Ensure path points to your fine-tuned model)
+MODEL_PATH = "path_to_fine_tuned_model"  # Replace with the correct model path
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
+
+# Set model to evaluation mode (Disables dropout for stable predictions)
+model.eval()
 
 # Define Label Mapping (Modify based on your dataset)
 LABEL_MAPPING = {
@@ -16,13 +19,17 @@ LABEL_MAPPING = {
     4: "Marketing Material"
 }
 
-# Classification Function
+# Optimized Classification Function
 def classify_text(text):
+    if not text.strip():
+        return "Please enter a valid business document text."
+
+    # Tokenize Input
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    
+
     with torch.no_grad():
         outputs = model(**inputs)
-    
+
     # Convert logits to probabilities
     probs = torch.nn.functional.softmax(outputs.logits, dim=1)
     
@@ -31,7 +38,11 @@ def classify_text(text):
     
     # Retrieve category name
     category = LABEL_MAPPING.get(label_idx, "Unknown")
-    
+
+    # Debugging Info (Uncomment for testing)
+    print(f"Logits: {outputs.logits}")
+    print(f"Probabilities: {probs}")
+
     return f"Predicted Category: {category} (Confidence: {probs[0][label_idx]:.2f})"
 
 # Gradio UI
@@ -39,7 +50,8 @@ demo = gr.Interface(
     fn=classify_text, 
     inputs=gr.Textbox(lines=4, placeholder="Enter business document text..."), 
     outputs="text", 
-    title="Multilingual Business Document Classifier"
+    title="Multilingual Business Document Classifier",
+    description="Classifies business documents into predefined categories using a multilingual model."
 )
 
 demo.launch()