Update model.py

model.py

@@ -1,7 +1,12 @@
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
+import logging
 from typing import Dict, List
 
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 def get_weather_condition(score: int) -> str:
     """Map weather impact score (0-100) to descriptive weather condition."""
     if score <= 10:
@@ -23,15 +28,18 @@ def call_ai_model_for_insights(input_data: Dict, delay_risk: float) -> List[str]
     """
     model_name = "sshleifer/distilbart-cnn-6-6"
     try:
-        # Load tokenizer and model for CPU
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False)
+        logger.info(f"Loading model: {model_name}")
+        # Load tokenizer and model with minimal memory usage
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False, use_fast=True)
         model = AutoModelForSeq2SeqLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float32,  # Use float32 for CPU
-            use_safetensors=True,  # Ensure safe loading
-            trust_remote_code=False
+            torch_dtype=torch.float32,  # CPU-compatible
+            use_safetensors=True,  # Secure loading
+            trust_remote_code=False,
+            low_cpu_mem_usage=True  # Optimize for low memory
         )
         
+        logger.info("Model loaded successfully. Generating insights...")
         # Prepare prompt
         prompt = f"""
         You are an AI assistant analyzing project delay risks for a construction project.
@@ -51,13 +59,13 @@ def call_ai_model_for_insights(input_data: Dict, delay_risk: float) -> List[str]
         Format the response as a list of strings, e.g., ["Insight 1", "Insight 2"].
         """
         
-        # Tokenize and generate with no_grad for memory efficiency
+        # Tokenize and generate with memory-efficient settings
         with torch.no_grad():
             inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to("cpu")
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=150,  # Smaller output for CPU efficiency
-                num_beams=4,  # Beam search for better quality
+                max_new_tokens=100,  # Reduced for faster CPU inference
+                num_beams=4,
                 temperature=0.7,
                 do_sample=True
             )
@@ -65,11 +73,22 @@ def call_ai_model_for_insights(input_data: Dict, delay_risk: float) -> List[str]
         
         # Parse response into a list
         insights = [line.strip() for line in response.split("\n") if line.strip() and line.strip() not in [prompt]]
-        return insights[:4]  # Limit to 2-4 insights
+        logger.info(f"Generated insights: {insights}")
+        return insights[:4] or ["No insights generated; review input data."]
 
     except Exception as e:
-        print(f"Error with model inference: {e}")
-        return ["AI model unavailable; monitor progress and resource allocation."]
+        logger.error(f"Model inference failed: {str(e)}")
+        # Fallback: Generate basic rule-based insights
+        fallback_insights = []
+        if delay_risk > 75:
+            fallback_insights.append("High risk detected; allocate additional resources urgently.")
+        elif delay_risk > 50:
+            fallback_insights.append("Moderate risk; consider extending shift hours or hiring staff.")
+        if input_data.get('workforce_gap', 0) > 20:
+            fallback_insights.append("Significant workforce gap; recruit additional workers.")
+        if input_data.get('weather_impact_score', 0) > 50:
+            fallback_insights.append("Adverse weather; prioritize indoor tasks.")
+        return fallback_insights or ["AI model unavailable; monitor progress and resource allocation."]
 
 def predict_delay(input_data: Dict) -> Dict:
     """
@@ -77,6 +96,7 @@ def predict_delay(input_data: Dict) -> Dict:
     Uses task duration, progress, workforce info, and weather impact.
     Insights are generated by DistilBART (CPU).
     """
+    logger.info("Starting delay prediction")
     phase = input_data.get("phase", "")
     task = input_data.get("task", "")
     expected_duration = input_data.get("task_expected_duration", 0)
@@ -153,6 +173,7 @@ def predict_delay(input_data: Dict) -> Dict:
     # Generate AI-driven insights
     insights = call_ai_model_for_insights(input_data, delay_risk)
 
+    logger.info(f"Prediction completed: Delay risk = {delay_risk:.1f}%")
     return {
         "project": input_data.get("project_name", "Unnamed Project"),
         "phase": phase,