Improve model loading with better error handling and SafeTensors support

app.py

@@ -12,81 +12,87 @@ import uvicorn
 # Initialize FastAPI
 app = FastAPI()
 
-# Download model files
+# Model configuration
 model_name = "Hayme/agrisago-bert"
-print("Downloading model files...")
+print("Loading AgriSagot BERT model...")
 
 try:
-    # Download model files
-    model_dir = "./model"
-    os.makedirs(model_dir, exist_ok=True)
+    print("Attempting to load tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    print("✅ Tokenizer loaded successfully")
     
-    # Download necessary files
-    files_to_download = [
-        "config.json",
-        "pytorch_model.bin", 
-        "tokenizer_config.json",
-        "tokenizer.json",
-        "vocab.txt"
-    ]
-    
-    for file in files_to_download:
-        try:
-            local_path = hf_hub_download(
-                repo_id=model_name,
-                filename=file,
-                local_dir=model_dir
-            )
-            print(f"Downloaded {file}")
-        except Exception as e:
-            print(f"Warning: Could not download {file}: {e}")
-    
-    # Load tokenizer and model
-    print("Loading tokenizer and model...")
-    tokenizer = AutoTokenizer.from_pretrained(model_dir)
-    model = AutoModel.from_pretrained(model_dir)
+    print("Attempting to load model...")
+    # Try to load model - transformers will automatically handle safetensors vs pytorch format
+    model = AutoModel.from_pretrained(
+        model_name, 
+        trust_remote_code=True,
+        torch_dtype=torch.float32,  # Ensure compatibility
+        device_map="auto" if torch.cuda.is_available() else None
+    )
     model.eval()
-    
-    print("Model loaded successfully!")
+    print("✅ Model loaded successfully")
     model_loaded = True
     
+    # Get model info
+    print(f"Model type: {type(model)}")
+    print(f"Model device: {next(model.parameters()).device}")
+    
 except Exception as e:
-    print(f"Error loading model: {e}")
-    print("Using fallback dummy model for testing")
+    print(f"❌ Error loading model: {e}")
+    print("Will use dummy responses for testing")
     tokenizer = None
     model = None
     model_loaded = False
 
 def get_bert_embedding(text):
     """Get BERT embedding for text"""
-    if not model_loaded:
+    if not model_loaded or model is None or tokenizer is None:
+        print("⚠️ Model not loaded, returning dummy embedding")
         # Return a dummy embedding for testing
-        return [0.1] * 768
+        np.random.seed(hash(text) % 2**32)  # Consistent dummy based on text
+        return np.random.normal(0, 0.1, 768).tolist()
     
     try:
+        print(f"Getting embedding for: {text[:50]}...")
+        
         # Tokenize and encode
-        inputs = tokenizer(text, 
-                          return_tensors="pt", 
-                          truncation=True, 
-                          padding=True, 
-                          max_length=512)
+        inputs = tokenizer(
+            text, 
+            return_tensors="pt", 
+            truncation=True, 
+            padding=True, 
+            max_length=512
+        )
+        
+        # Move to same device as model
+        device = next(model.parameters()).device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
         
         # Get embeddings
         with torch.no_grad():
             outputs = model(**inputs)
             # Use [CLS] token embedding (first token)
-            embedding = outputs.last_hidden_state[:, 0, :].squeeze().numpy()
+            embedding = outputs.last_hidden_state[:, 0, :].squeeze()
             
+            # Move back to CPU and convert to numpy
+            if embedding.device != torch.device('cpu'):
+                embedding = embedding.cpu()
+            embedding = embedding.numpy()
+            
+        print(f"✅ Generated embedding of shape: {embedding.shape}")
         return embedding.tolist()
         
     except Exception as e:
-        print(f"Error in get_bert_embedding: {e}")
+        print(f"❌ Error in get_bert_embedding: {e}")
         # Return dummy embedding on error
-        return [0.1] * 768
+        np.random.seed(hash(text) % 2**32)
+        return np.random.normal(0, 0.1, 768).tolist()
 
 def calculate_similarity(text1, text2):
     """Calculate cosine similarity between two texts"""
     try:
+        print(f"Calculating similarity between texts...")
+        
         # Get embeddings
         emb1 = np.array(get_bert_embedding(text1))
         emb2 = np.array(get_bert_embedding(text2))
@@ -100,10 +106,12 @@ def calculate_similarity(text1, text2):
             return 0.0
             
         similarity = dot_product / (norm1 * norm2)
-        return float(similarity)
+        result = float(similarity)
+        print(f"✅ Similarity calculated: {result:.4f}")
+        return result
         
     except Exception as e:
-        print(f"Error calculating similarity: {e}")
+        print(f"❌ Error calculating similarity: {e}")
         return 0.0
 
 # FastAPI endpoints
@@ -119,8 +127,14 @@ async def get_embedding_endpoint(request: TextRequest):
     """Get embedding for text"""
     try:
         embedding = get_bert_embedding(request.text)
-        return {"embedding": embedding, "success": True}
+        return {
+            "embedding": embedding, 
+            "success": True,
+            "model_loaded": model_loaded,
+            "embedding_length": len(embedding)
+        }
     except Exception as e:
+        print(f"API Error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/similarity")
@@ -128,15 +142,36 @@ async def get_similarity_endpoint(request: SimilarityRequest):
     """Get similarity between two texts"""
     try:
         similarity = calculate_similarity(request.text1, request.text2)
-        return {"similarity": similarity, "success": True}
+        return {
+            "similarity": similarity, 
+            "success": True,
+            "model_loaded": model_loaded
+        }
     except Exception as e:
+        print(f"API Error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
 async def root():
-    return {"message": "AgriSagot BERT Model API", "status": "running", "model_loaded": model_loaded}
+    return {
+        "message": "AgriSagot BERT Model API", 
+        "status": "running", 
+        "model_loaded": model_loaded,
+        "model_name": model_name,
+        "torch_version": torch.__version__,
+        "device": str(next(model.parameters()).device) if model_loaded else "N/A"
+    }
 
-# Gradio interface
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "model_loaded": model_loaded,
+        "endpoints": ["/", "/embedding", "/similarity", "/health"]
+    }
+
+# Gradio interface functions
 def gradio_embedding(text):
     """Gradio interface for embeddings"""
     if not text.strip():
@@ -144,7 +179,8 @@ def gradio_embedding(text):
     
     try:
         embedding = get_bert_embedding(text)
-        return f"Generated embedding vector of length {len(embedding)}\nFirst 10 values: {embedding[:10]}"
+        status = "✅ Real BERT embedding" if model_loaded else "⚠️ Dummy embedding (model not loaded)"
+        return f"{status}\nEmbedding length: {len(embedding)}\nFirst 10 values: {embedding[:10]}"
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -155,40 +191,97 @@ def gradio_similarity(text1, text2):
     
     try:
         similarity = calculate_similarity(text1, text2)
-        return f"Similarity score: {similarity:.4f}"
+        status = "✅ Real BERT similarity" if model_loaded else "⚠️ Dummy similarity (model not loaded)"
+        return f"{status}\nSimilarity score: {similarity:.4f}"
     except Exception as e:
         return f"Error: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="AgriSagot BERT Model") as demo:
-    gr.Markdown("# AgriSagot BERT Model")
-    gr.Markdown("Agricultural text processing with BERT embeddings")
+with gr.Blocks(title="AgriSagot BERT Model", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌾 AgriSagot BERT Model")
+    gr.Markdown("Agricultural text processing with BERT embeddings for crop disease recommendations")
     
-    if not model_loaded:
-        gr.Markdown("⚠️ **Warning**: Model not loaded properly. Using dummy responses for testing.")
+    # Status display
+    if model_loaded:
+        gr.Markdown("✅ **Status**: Model loaded successfully! Using real BERT embeddings.")
     else:
-        gr.Markdown("✅ Model loaded successfully!")
+        gr.Markdown("⚠️ **Status**: Model not loaded. Using dummy responses for API testing.")
+    
+    gr.Markdown(f"**Model**: {model_name}")
+    gr.Markdown(f"**PyTorch Version**: {torch.__version__}")
     
-    with gr.Tab("Text Embedding"):
-        text_input = gr.Textbox(label="Enter text", placeholder="e.g., 'Cabbage fungal treatment'")
-        embedding_output = gr.Textbox(label="Embedding Info", lines=3)
-        embedding_btn = gr.Button("Get Embedding")
+    with gr.Tab("🔍 Text Embedding"):
+        gr.Markdown("Generate BERT embeddings for agricultural text")
+        text_input = gr.Textbox(
+            label="Enter agricultural text", 
+            placeholder="e.g., 'Cabbage fungal treatment with copper-based fungicide'",
+            lines=2
+        )
+        embedding_output = gr.Textbox(label="Embedding Info", lines=4)
+        embedding_btn = gr.Button("Get Embedding", variant="primary")
         embedding_btn.click(gradio_embedding, inputs=text_input, outputs=embedding_output)
     
-    with gr.Tab("Text Similarity"):
-        text1_input = gr.Textbox(label="Text 1", placeholder="e.g., 'Cabbage disease treatment'")
-        text2_input = gr.Textbox(label="Text 2", placeholder="e.g., 'Fungicide for cabbage'")
-        similarity_output = gr.Textbox(label="Similarity Score")
-        similarity_btn = gr.Button("Calculate Similarity")
+    with gr.Tab("🔄 Text Similarity"):
+        gr.Markdown("Compare similarity between two agricultural texts")
+        text1_input = gr.Textbox(
+            label="Text 1", 
+            placeholder="e.g., 'Cabbage disease treatment'",
+            lines=2
+        )
+        text2_input = gr.Textbox(
+            label="Text 2", 
+            placeholder="e.g., 'Fungicide for cabbage crops'",
+            lines=2
+        )
+        similarity_output = gr.Textbox(label="Similarity Result", lines=3)
+        similarity_btn = gr.Button("Calculate Similarity", variant="primary")
         similarity_btn.click(gradio_similarity, inputs=[text1_input, text2_input], outputs=similarity_output)
+    
+    with gr.Tab("📚 API Documentation"):
+        gr.Markdown("""
+        ## API Endpoints
+        
+        ### POST /embedding
+        Get BERT embedding for text
+        ```json
+        {
+          "text": "your agricultural text here"
+        }
+        ```
+        
+        ### POST /similarity  
+        Get similarity between two texts
+        ```json
+        {
+          "text1": "first text",
+          "text2": "second text"
+        }
+        ```
+        
+        ### GET /health
+        Check API health status
+        
+        ## Example Usage
+        ```bash
+        curl -X POST "https://hayme-agrisagot-bert.hf.space/embedding" \\
+             -H "Content-Type: application/json" \\
+             -d '{"text":"cabbage fungal disease treatment"}'
+        ```
+        """)
 
 # Mount Gradio app to FastAPI
 app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
-    print("Starting server...")
-    print("FastAPI endpoints available at:")
+    print("\n" + "="*50)
+    print("🚀 Starting AgriSagot BERT API Server")
+    print("="*50)
+    print(f"Model loaded: {model_loaded}")
+    print("FastAPI endpoints:")
+    print("- GET  /")
+    print("- GET  /health") 
     print("- POST /embedding")
-    print("- POST /similarity") 
+    print("- POST /similarity")
     print("- Gradio interface at /")
+    print("="*50)
     uvicorn.run(app, host="0.0.0.0", port=7860)