Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from simple_salesforce import Salesforce
 import os
 from dotenv import load_dotenv
@@ -8,32 +8,28 @@ from dotenv import load_dotenv
 # Load environment variables
 load_dotenv()
 
-# Required env vars check
+# Check for required environment variables
 required_env_vars = ['SF_USERNAME', 'SF_PASSWORD', 'SF_SECURITY_TOKEN']
 missing_vars = [var for var in required_env_vars if not os.getenv(var)]
 if missing_vars:
     raise EnvironmentError(f"Missing required environment variables: {missing_vars}")
 
-# Defaults
+# Configurable defaults
 KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'
 ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))
 
-# Load model and tokenizer
-model_name = "distilgpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
+# Use instruction-tuned model
+model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
-model.config.pad_token_id = tokenizer.pad_token_id
+# Prompt template
+PROMPT_TEMPLATE = """
+You are an AI coach for construction supervisors. Based on the information below, generate:
 
-# Better Prompt
-PROMPT_TEMPLATE = """You are an AI coach for construction supervisors. Given the role, project, milestones, and a reflection log, generate:
-
-1. A Daily Checklist based on the milestones and role.
-2. Focus Suggestions based on concerns or keywords in the reflection log.
-3. A Motivational Quote for the day.
+1. A Daily Checklist based on role and milestones.
+2. Focus Suggestions based on reflection.
+3. A Motivational Quote.
 
 Inputs:
 Role: {role}
@@ -43,15 +39,15 @@ Reflection Log: {reflection}
 
 Output Format:
 Checklist:
-- 
+-
 
 Suggestions:
-- 
+-
 
 Quote:
 """
 
-# Salesforce Functions
+# Salesforce functions
 def get_roles_from_salesforce():
     try:
         sf = Salesforce(
@@ -74,7 +70,6 @@ def get_supervisor_name_by_role(role):
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-        role = role.replace("'", "\\'")
         result = sf.query(f"SELECT Name FROM Supervisor__c WHERE Role__c = '{role}'")
         return [record['Name'] for record in result['records']]
     except Exception as e:
@@ -89,23 +84,22 @@ def get_projects_for_supervisor(supervisor_name):
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-        supervisor_name = supervisor_name.replace("'", "\\'")
-        supervisor_result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
-        if supervisor_result['totalSize'] == 0:
+        result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
+        if result['totalSize'] == 0:
             return ""
-        supervisor_id = supervisor_result['records'][0]['Id']
-        project_result = sf.query(f"SELECT Name FROM Project__c WHERE Supervisor_ID__c = '{supervisor_id}' LIMIT 1")
-        return project_result['records'][0]['Name'] if project_result['totalSize'] > 0 else ""
+        supervisor_id = result['records'][0]['Id']
+        project = sf.query(f"SELECT Name FROM Project__c WHERE Supervisor_ID__c = '{supervisor_id}' LIMIT 1")
+        return project['records'][0]['Name'] if project['totalSize'] else ""
     except Exception as e:
         print(f"⚠️ Error fetching project: {e}")
         return ""
 
 def field_exists(sf, object_name, field_name):
     try:
-        obj_desc = getattr(sf, object_name).describe()
-        return field_name in [field['name'] for field in obj_desc['fields']]
+        desc = getattr(sf, object_name).describe()
+        return field_name in [f['name'] for f in desc['fields']]
     except Exception as e:
-        print(f"⚠️ Error checking field {field_name}: {e}")
+        print(f"⚠️ Error checking field: {e}")
         return False
 
 def save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_name):
@@ -116,39 +110,37 @@ def save_to_salesforce(role, project_id, milestones, reflection, checklist, sugg
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-        supervisor_result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
-        if supervisor_result['totalSize'] == 0:
-            print("❌ No supervisor found.")
+        supervisor = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
+        if not supervisor['totalSize']:
             return
-        supervisor_id = supervisor_result['records'][0]['Id']
-        project_result = sf.query(f"SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
-        if project_result['totalSize'] == 0:
-            print("❌ No project found.")
+        supervisor_id = supervisor['records'][0]['Id']
+
+        project = sf.query(f"SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
+        if not project['totalSize']:
             return
-        project_record_id = project_result['records'][0]['Id']
+        project_id_val = project['records'][0]['Id']
 
         data = {
             'Supervisor_ID__c': supervisor_id,
-            'Project_ID__c': project_record_id,
+            'Project_ID__c': project_id_val,
             'Daily_Checklist__c': checklist[:255],
             'Suggested_Tips__c': suggestions[:255],
             'Reflection_Log__c': reflection[:255],
             'KPI_Flag__c': KPI_FLAG_DEFAULT,
             'Engagement_Score__c': ENGAGEMENT_SCORE_DEFAULT
         }
-
         if field_exists(sf, 'Supervisor_AI_Coaching__c', 'Milestones_KPIs__c'):
             data['Milestones_KPIs__c'] = milestones[:255]
 
         sf.Supervisor_AI_Coaching__c.create(data)
-        print("✅ Record created in Salesforce.")
+        print("✅ Record saved.")
     except Exception as e:
         print(f"❌ Error saving to Salesforce: {e}")
 
-# Generate Function
+# Generation logic
 def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
     if not all([role, supervisor_name, project_id, milestones, reflection]):
-        return "❗ Please fill all fields.", "", ""
+        return "Please fill all fields.", "", ""
 
     prompt = PROMPT_TEMPLATE.format(
         role=role,
@@ -157,48 +149,49 @@ def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
         reflection=reflection
     )
 
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     try:
         with torch.no_grad():
-            outputs = model.generate(
+            output_ids = model.generate(
                 inputs['input_ids'],
-                max_new_tokens=150,
-                no_repeat_ngram_size=2,
+                max_new_tokens=256,
                 do_sample=True,
                 top_p=0.9,
-                temperature=0.8,
+                temperature=0.7,
                 pad_token_id=tokenizer.pad_token_id
             )
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     except Exception as e:
-        print(f"⚠️ Generation error: {e}")
-        return "Error during text generation", "", ""
+        return f"Error: {e}", "", ""
 
-    def extract_between(text, start, end):
+    def extract(text, start, end=None):
         s = text.find(start)
+        if s == -1:
+            return ""
+        s += len(start)
         e = text.find(end, s) if end else len(text)
-        return text[s + len(start):e].strip() if s != -1 else "Not found"
+        return text[s:e].strip()
 
-    checklist = extract_between(result, "Checklist:\n", "Suggestions:")
-    suggestions = extract_between(result, "Suggestions:\n", "Quote:")
-    quote = extract_between(result, "Quote:\n", None)
+    checklist = extract(result, "Checklist:\n", "Suggestions:")
+    suggestions = extract(result, "Suggestions:\n", "Quote:")
+    quote = extract(result, "Quote:\n")
 
     save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_name)
     return checklist, suggestions, quote
 
-# Interface
+# Gradio UI
 def create_interface():
     roles = get_roles_from_salesforce()
     with gr.Blocks(theme="soft") as demo:
-        gr.Markdown("## 🧠 AI-Powered Supervisor Assistant")
+        gr.Markdown("## 👷‍♂️ AI Coaching Assistant for Construction Supervisors")
 
         with gr.Row():
             role = gr.Dropdown(choices=roles, label="Role")
             supervisor_name = gr.Dropdown(choices=[], label="Supervisor Name")
             project_id = gr.Textbox(label="Project ID", interactive=False)
 
-        milestones = gr.Textbox(label="Milestones (comma-separated KPIs)", placeholder="E.g. Safety training, daily inspection")
-        reflection = gr.Textbox(label="Reflection Log", lines=4, placeholder="Any concerns, delays, updates...")
+        milestones = gr.Textbox(label="Milestones (comma-separated)", placeholder="e.g., Safety check, Equipment inspection")
+        reflection = gr.Textbox(label="Reflection Log", lines=4, placeholder="Any thoughts, challenges, delays...")
 
         with gr.Row():
             generate = gr.Button("Generate")
@@ -216,9 +209,7 @@ def create_interface():
                        inputs=[role, supervisor_name, project_id, milestones, reflection],
                        outputs=[checklist_output, suggestions_output, quote_output])
 
-        clear.click(fn=lambda: ("", "", "", "", ""), inputs=None,
-                    outputs=[role, supervisor_name, project_id, milestones, reflection])
-
+        clear.click(fn=lambda: ("", "", "", "", ""), outputs=[role, supervisor_name, project_id, milestones, reflection])
         refresh.click(fn=lambda: gr.update(choices=get_roles_from_salesforce()), outputs=role)
 
     return demo