Update app.py

app.py

@@ -4,6 +4,11 @@ from transformers import AutoTokenizer, AlbertForSequenceClassification
 import numpy as np
 import os
 import gdown
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Define Google Drive folder IDs for each model
 model_drive_ids = {
@@ -21,11 +26,16 @@ os.makedirs(save_dir, exist_ok=True)
 for task, folder_id in model_drive_ids.items():
     output_dir = os.path.join(save_dir, task)
     if not os.path.exists(output_dir):
-        gdown.download_folder(
-            f"https://drive.google.com/drive/folders/1kEXKoJxxD5-0FO8WvtagzseSIC5q-rRY?usp=sharing/{folder_id}",
-            output=output_dir,
-            quiet=False
-        )
+        logger.info(f"Downloading {task} model from Google Drive...")
+        try:
+            gdown.download_folder(
+                f"https://drive.google.com/drive/folders/1kEXKoJxxD5-0FO8WvtagzseSIC5q-rRY?usp=sharing/{folder_id}",
+                output=output_dir,
+                quiet=False
+            )
+        except Exception as e:
+            logger.error(f"Failed to download {task} model: {str(e)}")
+            raise
 
 # Define model paths
 tasks = ["sentiment", "emotion", "hate_speech", "sarcasm"]
@@ -39,33 +49,48 @@ label_mappings = {
     "sarcasm": ["no", "yes"]
 }
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indic-bert")
+# Load tokenizer with use_fast=False to avoid fast tokenizer issues
+try:
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indic-bert", use_fast=False)
+except Exception as e:
+    logger.error(f"Failed to load tokenizer: {str(e)}")
+    raise
 
 # Load all models
 models = {}
 for task in tasks:
-    if not os.path.exists(model_paths[task]):
-        raise FileNotFoundError(f"Model directory {model_paths[task]} not found.")
-    models[task] = AlbertForSequenceClassification.from_pretrained(model_paths[task])
+    model_path = model_paths[task]
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model directory {model_path} not found.")
+    try:
+        logger.info(f"Loading {task} model...")
+        models[task] = AlbertForSequenceClassification.from_pretrained(model_path)
+    except Exception as e:
+        logger.error(f"Failed to load {task} model: {str(e)}")
+        raise
 
 # Function to predict for a single task
 def predict_task(text, task, model, tokenizer, max_length=128):
-    inputs = tokenizer(
-        text,
-        padding=True,
-        truncation=True,
-        max_length=max_length,
-        return_tensors="pt"
-    )
-    
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
-    
-    labels = label_mappings[task]
-    return {label: f"{prob*100:.2f}%" for label, prob in zip(labels, probabilities)}
+    try:
+        inputs = tokenizer(
+            text,
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors="pt"
+        )
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            probabilities = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
+        
+        labels = label_mappings[task]
+        return {label: f"{prob*100:.2f}%" for label, prob in zip(labels, probabilities)}
+    except Exception as e:
+        logger.error(f"Error predicting for {task}: {str(e)}")
+        return {label: "Error" for label in label_mappings[task]}
 
 # Gradio interface function
 def predict_all_tasks(text):
@@ -94,4 +119,5 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
+    logger.info("Launching Gradio interface...")
     iface.launch(server_name="0.0.0.0", server_port=7860)