Update model.py

model.py

@@ -25,62 +25,56 @@ def get_weather_condition(score: int) -> str:
 
 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.
+    Use T5-Small in Hugging Face Space (CPU) to generate insights based on input data and delay risk.
     """
-    model_name = "sshleifer/distilbart-cnn-6-6"
+    model_name = "t5-small"
     max_retries = 3
     retry_delay = 5  # seconds
     
     for attempt in range(max_retries):
         try:
             logger.info(f"Attempt {attempt + 1}/{max_retries} - 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,  # CPU-compatible
-                use_safetensors=True,  # Secure loading
+                torch_dtype=torch.float32,
+                use_safetensors=True,
                 trust_remote_code=False,
-                low_cpu_mem_usage=True  # Optimize for low memory
+                low_cpu_mem_usage=True
             )
             
             logger.info("Model loaded successfully. Generating insights...")
-            # 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"].
+            Summarize the following project delay risk data into 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, e.g., ["Insight 1", "Insight 2"].
             """
             
-            # 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=100,  # Reduced for faster CPU inference
+                    max_new_tokens=100,
                     num_beams=4,
                     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]]
             logger.info(f"Generated insights: {insights}")
             return insights[:4] or ["No insights generated; review input data."]
-
         except Exception as e:
             logger.error(f"Attempt {attempt + 1}/{max_retries} - Model inference failed: {str(e)}")
             if attempt < max_retries - 1:
@@ -88,7 +82,6 @@ def call_ai_model_for_insights(input_data: Dict, delay_risk: float) -> List[str]
                 time.sleep(retry_delay)
             else:
                 logger.error("Max retries reached. Using fallback insights.")
-                # Fallback: Generate basic rule-based insights
                 fallback_insights = []
                 if delay_risk > 75:
                     fallback_insights.append("High risk detected; allocate additional resources urgently.")
@@ -98,13 +91,13 @@ def call_ai_model_for_insights(input_data: Dict, delay_risk: float) -> List[str]
                     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."]
+                return fallback_insights or ["AI model failed to generate insights; check system resources."]
 
 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).
+    Insights are generated by T5-Small (CPU).
     """
     logger.info("Starting delay prediction")
     phase = input_data.get("phase", "")