Update model.py

model.py

@@ -1,4 +1,8 @@
-def get_weather_condition(score):
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
+from typing import Dict, List
+
+def get_weather_condition(score: int) -> str:
     """Map weather impact score (0-100) to descriptive weather condition."""
     if score <= 10:
         return "Sunny"
@@ -13,10 +17,65 @@ def get_weather_condition(score):
     else:
         return "Severe Storm"
 
-def predict_delay(input_data):
+def call_ai_model_for_insights(input_data: Dict, delay_risk: float) -> List[str]:
+    """
+    Use DistilBART in Hugging Face Space (CPU) to generate insights based on input data and delay risk.
+    """
+    model_name = "sshleifer/distilbart-cnn-6-6"
+    try:
+        # Load tokenizer and model for CPU
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False)
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,  # Use float32 for CPU
+            use_safetensors=True,  # Ensure safe loading
+            trust_remote_code=False
+        )
+        
+        # Prepare prompt
+        prompt = f"""
+        You are an AI assistant analyzing project delay risks for a construction project.
+        Based on the following data, provide 2-4 concise insights or mitigation strategies as a list:
+        - Project: {input_data.get('project_name', 'Unnamed Project')}
+        - Phase: {input_data.get('phase', '')}
+        - Task: {input_data.get('task', '')}
+        - Expected Duration: {input_data.get('task_expected_duration', 0)} days
+        - Actual Duration: {input_data.get('task_actual_duration', 0)} days
+        - Current Progress: {input_data.get('current_progress', 0)}%
+        - Workforce Gap: {input_data.get('workforce_gap', 0)}%
+        - Workforce Skill Level: {input_data.get('workforce_skill_level', '').lower()}
+        - Shift Hours: {input_data.get('workforce_shift_hours', 0)} hours
+        - Weather Impact Score: {input_data.get('weather_impact_score', 0)} (Condition: {get_weather_condition(input_data.get('weather_impact_score', 0))})
+        - Calculated Delay Risk: {delay_risk:.1f}%
+
+        Format the response as a list of strings, e.g., ["Insight 1", "Insight 2"].
+        """
+        
+        # Tokenize and generate with no_grad for memory efficiency
+        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
+                temperature=0.7,
+                do_sample=True
+            )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # 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
+
+    except Exception as e:
+        print(f"Error with model inference: {e}")
+        return ["AI model unavailable; monitor progress and resource allocation."]
+
+def predict_delay(input_data: Dict) -> Dict:
     """
     Predict delay probability based on project task data.
     Uses task duration, progress, workforce info, and weather impact.
+    Insights are generated by DistilBART (CPU).
     """
     phase = input_data.get("phase", "")
     task = input_data.get("task", "")
@@ -42,13 +101,11 @@ def predict_delay(input_data):
     }
 
     delay_risk = 0
-    insights = []
 
     # 1. Duration overrun risk
     if expected_duration > 0 and actual_duration > expected_duration:
         overrun_pct = ((actual_duration - expected_duration) / expected_duration) * 100
         delay_risk += min(overrun_pct, 30)
-        insights.append(f"Actual duration is {overrun_pct:.1f}% over expected.")
 
     # 2. Progress lag risk
     if expected_duration > 0 and current_progress >= 0:
@@ -56,36 +113,26 @@ def predict_delay(input_data):
         if current_progress < expected_progress:
             progress_gap = expected_progress - current_progress
             delay_risk += min(progress_gap, 25)
-            insights.append(f"Current progress ({current_progress}%) lags behind expected ({expected_progress:.1f}%).")
 
     # 3. Workforce gap impact
     if workforce_gap_pct > 0:
         delay_risk += min(workforce_gap_pct * 0.5, 20)
-        insights.append(f"Workforce gap at {workforce_gap_pct}% reduces productivity.")
 
     # 4. Skill level effect
     if skill_level == "low":
         delay_risk += 15
-        insights.append("Low skill level may reduce task efficiency.")
     elif skill_level == "medium":
         delay_risk += 7
 
     # 5. Shift hours effect
     if shift_hours < 8:
-        penalty = (8 - shift_hours) * 3
-        delay_risk += penalty
-        insights.append(f"Reduced shift hours ({shift_hours}h) add {penalty:.1f}% to delay risk.")
+        delay_risk += (8 - shift_hours) * 3
     elif shift_hours > 8:
-        bonus = (shift_hours - 8) * 2
-        delay_risk -= min(bonus, 10)
-        insights.append(f"Extended shift hours ({shift_hours}h) reduce delay risk by {min(bonus, 10):.1f}%.")
-    else:
-        insights.append("Standard shift hours (8h) have neutral impact.")
+        delay_risk -= min((shift_hours - 8) * 2, 10)
 
     # 6. Weather impact effect
     if weather_score > 50:
         delay_risk += min(weather_score / 2, 20)
-        insights.append(f"High weather impact score ({weather_score}) — current condition: {weather_condition}.")
 
     # Ensure delay_risk is between 0 and 100
     delay_risk = max(0, min(delay_risk, 100))
@@ -95,7 +142,6 @@ def predict_delay(input_data):
     if phase in task_options:
         for t in task_options[phase]:
             task_risk = delay_risk
-            # Adjust risk slightly for other tasks (simulate variation)
             if t != task:
                 task_risk = min(max(task_risk + (hash(t) % 10 - 5), 0), 100)  # ±5% variation
             high_risk_phases.append({
@@ -104,21 +150,8 @@ def predict_delay(input_data):
                 "risk": round(task_risk, 1)
             })
 
-    # Enhanced mitigation strategies based on risk level
-    mitigation_strategies = []
-    if delay_risk > 75:
-        mitigation_strategies.append("Urgent: Allocate additional resources and expedite critical tasks.")
-        if phase == "Construction":
-            mitigation_strategies.append("Secure backup equipment and materials to counter weather delays.")
-        elif phase == "Planning":
-            mitigation_strategies.append("Fast-track permit approvals and stakeholder alignment.")
-    elif delay_risk > 50:
-        mitigation_strategies.append("Moderate risk: Increase workforce or extend shift hours.")
-        if weather_score > 50:
-            mitigation_strategies.append("Plan indoor tasks or weather-resistant schedules.")
-    else:
-        mitigation_strategies.append("Low risk: Maintain current plan, monitor progress closely.")
-    insights.extend(mitigation_strategies)
+    # Generate AI-driven insights
+    insights = call_ai_model_for_insights(input_data, delay_risk)
 
     return {
         "project": input_data.get("project_name", "Unnamed Project"),