Update app.py

app.py

@@ -3,17 +3,25 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from joblib import Memory
 import datetime
+import os
+
+# Enable offline mode for Hugging Face transformers
+os.environ["HF_HUB_OFFLINE"] = "1"
 
 # Initialize cache
 cache_dir = "./cache"
 memory = Memory(cache_dir, verbose=0)
 
-# Load pre-trained model and tokenizer (using distilgpt2 for testing)
+# Load pre-trained model and tokenizer with local_files_only=True
 model_name = "distilgpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, local_files_only=True)
+
+# Set pad_token_id to eos_token_id to avoid warnings
+tokenizer.pad_token = tokenizer.eos_token
+model.config.pad_token_id = tokenizer.eos_token_id
 
-# Define a more explicit prompt template
+# Define a more explicit prompt template with few-shot examples
 PROMPT_TEMPLATE = """You are an AI coach for construction supervisors. Based on the following inputs, generate a daily checklist, focus suggestions, and a motivational quote. Format your response with clear labels as follows:
 
 Checklist:
@@ -25,13 +33,47 @@ Suggestions:
 Quote:
 - Your motivational quote here
 
+Here are examples of expected outputs:
+
+Example 1:
+Inputs:
+Role: Supervisor
+Project: PROJ-01
+Milestones: Lay foundation, Check materials
+Reflection: Team is on track, but weather might delay work.
+Checklist:
+- Confirm foundation laying schedule by 8 AM.
+- Inspect materials for quality at 10 AM.
+Suggestions:
+- Monitor weather updates and plan contingencies.
+- Ensure team has protective gear for rain.
+Quote:
+- Stay steady—every foundation leads to greater heights!
+
+Example 2:
+Inputs:
+Role: Foreman
+Project: PROJ-02
+Milestones: Install beams, Safety check
+Reflection: Minor delay due to late delivery.
+Checklist:
+- Oversee beam installation starting at 9 AM.
+- Conduct safety check at 1 PM.
+Suggestions:
+- Follow up with supplier to prevent future delays.
+- Review safety protocols with the team.
+Quote:
+- Safety first, progress always—keep the momentum!
+
+Now, generate the checklist, suggestions, and quote for the following inputs:
+
 Inputs:
 Role: {role}
 Project: {project_id}
 Milestones: {milestones}
 Reflection: {reflection}
 
-Now, generate the checklist, suggestions, and quote."""
+Generate the response in the same format as the examples above."""
 
 # Cache reset check
 last_reset = datetime.date.today()
@@ -43,7 +85,7 @@ def reset_cache_if_new_day():
         memory.clear()
         last_reset = today
 
-# Cached generation function with improved parsing
+# Cached generation function with improved parsing and context-aware fallbacks
 @memory.cache
 def generate_outputs(role, project_id, milestones, reflection):
     reset_cache_if_new_day()
@@ -60,16 +102,27 @@ def generate_outputs(role, project_id, milestones, reflection):
         reflection=reflection
     )
     
-    # Tokenize and generate
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+    # Tokenize with attention_mask
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        max_length=512,
+        truncation=True,
+        padding=True,
+        return_attention_mask=True
+    )
+    
+    # Generate with attention_mask
     outputs = model.generate(
         inputs["input_ids"],
-        max_length=1500,  # Increased for longer outputs
+        attention_mask=inputs["attention_mask"],
+        max_length=1500,
         num_return_sequences=1,
         no_repeat_ngram_size=2,
         do_sample=True,
         top_p=0.9,
-        temperature=0.8  # Slightly higher for more creativity
+        temperature=0.8,
+        pad_token_id=tokenizer.eos_token_id
     )
     
     # Decode generated text
@@ -105,17 +158,30 @@ def generate_outputs(role, project_id, milestones, reflection):
         if not quote:
             quote = "No quote generated."
     
-    # Fallback if sections are still empty (due to model not being fine-tuned)
+    # Context-aware fallbacks based on inputs
     if checklist == "No checklist generated.":
-        checklist = "- Review milestones.\n- Assign tasks to team."
+        checklist_items = []
+        milestone_list = [m.strip() for m in milestones.split(",")]
+        for i, milestone in enumerate(milestone_list, 1):
+            checklist_items.append(f"- {milestone} by {8 + i*2} AM.")
+        checklist_items.append("- Check equipment status before end of day.")
+        checklist = "\n".join(checklist_items)
+    
     if suggestions == "No suggestions generated.":
-        suggestions = "- Coordinate with suppliers.\n- Plan for contingencies."
+        suggestions_items = []
+        if "equipment issues" in reflection.lower():
+            suggestions_items.append("- Schedule equipment maintenance to avoid future delays.")
+        if "suppliers" in reflection.lower():
+            suggestions_items.append("- Set up a morning call with suppliers to confirm timelines.")
+        suggestions_items.append("- Brief the team on tomorrow’s goals during the daily huddle.")
+        suggestions = "\n".join(suggestions_items if suggestions_items else ["- Coordinate with the team.", "- Plan for contingencies."])
+    
     if quote == "No quote generated.":
-        quote = "- Keep pushing forward!"
+        quote = "- Keep building—every step forward counts!"
     
     return checklist, suggestions, quote
 
-# Gradio interface (unchanged from original)
+# Gradio interface
 def create_interface():
     with gr.Blocks() as demo:
         gr.Markdown("# Construction Supervisor AI Coach")