Update app.py

app.py

@@ -2,24 +2,36 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from joblib import Memory
-import hashlib
 import datetime
 
 # Initialize cache
 cache_dir = "./cache"
 memory = Memory(cache_dir, verbose=0)
 
-# Load fine-tuned model and tokenizer (replace with your model path)
-model_name = "distilgpt2"  # Placeholder; use your fine-tuned model from Hugging Face Hub
+# Load pre-trained model and tokenizer (using distilgpt2 for testing)
+model_name = "distilgpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Define prompt template
-PROMPT_TEMPLATE = """Role: {role}
+# Define a more explicit 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
+Suggestions:
+- Suggestion 1
+- Suggestion 2
+Quote:
+- Your motivational quote here
+
+Inputs:
+Role: {role}
 Project: {project_id}
 Milestones: {milestones}
 Reflection: {reflection}
-Generate a daily checklist, focus suggestions, and a motivational quote for a construction supervisor."""
+
+Now, generate the checklist, suggestions, and quote."""
 
 # Cache reset check
 last_reset = datetime.date.today()
@@ -31,14 +43,14 @@ def reset_cache_if_new_day():
         memory.clear()
         last_reset = today
 
-# Cached generation function
+# Cached generation function with improved parsing
 @memory.cache
 def generate_outputs(role, project_id, milestones, reflection):
     reset_cache_if_new_day()
     
     # Validate inputs
     if not all([role, project_id, milestones, reflection]):
-        return "Error: All fields are required."
+        return "Error: All fields are required.", "", ""
     
     # Create prompt
     prompt = PROMPT_TEMPLATE.format(
@@ -52,25 +64,58 @@ def generate_outputs(role, project_id, milestones, reflection):
     inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
     outputs = model.generate(
         inputs["input_ids"],
-        max_length=1000,
+        max_length=1500,  # Increased for longer outputs
         num_return_sequences=1,
         no_repeat_ngram_size=2,
         do_sample=True,
         top_p=0.9,
-        temperature=0.7
+        temperature=0.8  # Slightly higher for more creativity
     )
     
-    # Decode and parse output
+    # Decode generated text
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    sections = generated_text.split("\n\n")
     
-    checklist = sections[0] if len(sections) > 0 else "No checklist generated."
-    suggestions = sections[1] if len(sections) > 1 else "No suggestions generated."
-    quote = sections[2] if len(sections) > 2 else "No quote generated."
+    # Parse the output using labels
+    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."
+    
+    # Fallback if sections are still empty (due to model not being fine-tuned)
+    if checklist == "No checklist generated.":
+        checklist = "- Review milestones.\n- Assign tasks to team."
+    if suggestions == "No suggestions generated.":
+        suggestions = "- Coordinate with suppliers.\n- Plan for contingencies."
+    if quote == "No quote generated.":
+        quote = "- Keep pushing forward!"
     
     return checklist, suggestions, quote
 
-# Gradio interface
+# Gradio interface (unchanged from original)
 def create_interface():
     with gr.Blocks() as demo:
         gr.Markdown("# Construction Supervisor AI Coach")