Update app.py

app.py

@@ -1,14 +1,9 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from joblib import Memory
 import datetime
 
-# Initialize cache
-cache_dir = "./cache"
-memory = Memory(cache_dir, verbose=0)
-
-# Load pre-trained model and tokenizer (allow online download)
+# Initialize model and tokenizer (preloading them for quicker response)
 model_name = "distilgpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
@@ -17,22 +12,18 @@ model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer.pad_token = tokenizer.eos_token
 model.config.pad_token_id = tokenizer.eos_token_id
 
-# Define a prompt template (structured format)
+# Define a more contextual prompt template
 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:
-- Item 1
-- Item 2
+- {milestones_list}
 
 Suggestions:
-- Suggestion 1
-- Suggestion 2
+- {suggestions_list}
 
 Quote:
 - Your motivational quote here
 
-Now, generate the checklist, suggestions, and quote for the following inputs:
-
 Inputs:
 Role: {role}
 Project: {project_id}
@@ -40,113 +31,75 @@ Milestones: {milestones}
 Reflection: {reflection}
 """
 
-# Cache reset check
-last_reset = datetime.date.today()
-
-def reset_cache_if_new_day():
-    global last_reset
-    today = datetime.date.today()
-    if today > last_reset:
-        memory.clear()
-        last_reset = today
-
-# Cached generation function with improved parsing and context-aware fallbacks
-@memory.cache
+# Function to generate outputs based on inputs
 def generate_outputs(role, project_id, milestones, reflection):
-    reset_cache_if_new_day()
-    
-    # Validate inputs
+    # Validate inputs to ensure no missing fields
     if not all([role, project_id, milestones, reflection]):
         return "Error: All fields are required.", "", ""
     
-    # Create prompt
+    # Create prompt from template
+    milestones_list = "\n- ".join([m.strip() for m in milestones.split(",")])
+    
+    suggestions_list = ""
+    if "delays" in reflection.lower():
+        suggestions_list = "- Consider adjusting timelines to accommodate delays.\n- Communicate delays to all relevant stakeholders."
+    elif "weather" in reflection.lower():
+        suggestions_list = "- Ensure team has rain gear.\n- Monitor weather updates for possible further delays."
+    elif "equipment" in reflection.lower():
+        suggestions_list = "- Inspect all equipment to ensure no malfunctions.\n- Schedule maintenance if necessary."
+
+    # Create final prompt
     prompt = PROMPT_TEMPLATE.format(
         role=role,
         project_id=project_id,
         milestones=milestones,
-        reflection=reflection
-    )
-    
-    # Tokenize with attention_mask
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt",
-        max_length=512,
-        truncation=True,
-        padding=True,
-        return_attention_mask=True
+        reflection=reflection,
+        milestones_list=milestones_list,
+        suggestions_list=suggestions_list
     )
+
+    # Tokenize inputs for model processing
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True, padding=True)
     
-    # Generate with attention_mask
-    outputs = model.generate(
-        inputs["input_ids"],
-        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,
-        pad_token_id=tokenizer.eos_token_id
-    )
+    # Generate response from the model
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs['input_ids'],
+            max_length=512,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.8,
+            pad_token_id=tokenizer.eos_token_id
+        )
     
-    # Decode generated text
+    # Decode the response
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Parse the output using labels
+
+    # Parse the output and ensure it is structured
     checklist = "No checklist generated."
     suggestions = "No suggestions generated."
     quote = "No quote generated."
-    
-    # Look for sections using labels
+
     if "Checklist:" in generated_text:
         checklist_start = generated_text.find("Checklist:") + len("Checklist:")
         suggestions_start = generated_text.find("Suggestions:")
-        if suggestions_start == -1:
-            suggestions_start = len(generated_text)
         checklist = generated_text[checklist_start:suggestions_start].strip()
-        if not checklist:
-            checklist = "No checklist generated."
     
     if "Suggestions:" in generated_text:
         suggestions_start = generated_text.find("Suggestions:") + len("Suggestions:")
         quote_start = generated_text.find("Quote:")
-        if quote_start == -1:
-            quote_start = len(generated_text)
         suggestions = generated_text[suggestions_start:quote_start].strip()
-        if not suggestions:
-            suggestions = "No suggestions generated."
     
     if "Quote:" in generated_text:
         quote_start = generated_text.find("Quote:") + len("Quote:")
         quote = generated_text[quote_start:].strip()
-        if not quote:
-            quote = "No quote generated."
-    
-    # Context-aware fallbacks based on inputs
-    if checklist == "No checklist generated.":
-        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_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 building—every step forward counts!"
-    
+
+    # Return structured outputs
     return checklist, suggestions, quote
 
-# Gradio interface
+# Gradio interface for fast user interaction
 def create_interface():
     with gr.Blocks() as demo:
         gr.Markdown("# Construction Supervisor AI Coach")