Update app.py

app.py

@@ -9,58 +9,73 @@ import os
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from peft import PeftModel
 
-# Configuration - Using better base models
+# Configuration - Multiple CPU-optimized models
 MODEL_PATHS = [
     "FrAnKu34t23/Construction_Risk_Prediction_Model_v3"
 ]
-# Better base model options - choose one based on your needs
-BASE_MODEL_ID = "microsoft/DialoGPT-medium"  # Better conversational model
-# Alternative options:
-# BASE_MODEL_ID = "gpt2-medium"  # Larger GPT-2
-# BASE_MODEL_ID = "microsoft/DialoGPT-large"  # Even better but slower
+
+# CPU-friendly base model options
+BASE_MODEL_ID = "microsoft/phi-2"  # Best balance of efficiency and capability
+# Alternative: "google/gemma-2b" for even better quality if your CPU can handle it
 
 models = []
 tokenizers = []
+model_names = []
 
-# Initialize better models for analysis
-injury_classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
-
-# Use a more capable model for text analysis and reasoning
-analysis_model = pipeline(
-    "text-generation", 
-    model="microsoft/DialoGPT-large",  # Better reasoning capabilities
-    device=0 if torch.cuda.is_available() else -1
+# CPU-friendly classifier
+injury_classifier = pipeline(
+    "zero-shot-classification", 
+    model="typeform/distilbert-base-uncased-mnli",
+    device=-1
 )
 
 def classify_injury_zero_shot(description):
     candidate_labels = [
         "Low severity injury (minor discomfort or bruise) or unrelevant cases",
-        "Medium severity injury (sprain, strain, moderate pain)",
+        "Medium severity injury (sprain, strain, moderate pain)", 
         "High severity injury (fracture, major trauma, amputation, fatal)"
     ]
     label_mapping = {
         candidate_labels[0]: "Low",
-        candidate_labels[1]: "Medium",
+        candidate_labels[1]: "Medium", 
         candidate_labels[2]: "High"
     }
     result = injury_classifier(description, candidate_labels)
     return label_mapping[result['labels'][0]]
 
 def load_models():
-    global models, tokenizers
+    global models, tokenizers, model_names
     try:
-        for path in MODEL_PATHS:
+        for i, path in enumerate(MODEL_PATHS):
+            print(f"Loading model {i+1}/{len(MODEL_PATHS)}: {path}")
+            
             tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
             tokenizer.pad_token = tokenizer.eos_token
-            base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID)
+            
+            # Load base model with CPU optimization
+            base_model = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL_ID,
+                torch_dtype=torch.float32,
+                device_map=None,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True  # Reduce memory usage
+            )
+            
+            # Load PEFT model
             model = PeftModel.from_pretrained(base_model, path)
             model.eval()
+            model = model.to('cpu')
+            
             models.append(model)
             tokenizers.append(tokenizer)
-        print("✅ All models loaded.")
+            model_names.append(f"Model_{i+1}")
+            
+        print(f"✅ All {len(models)} models loaded successfully on CPU.")
         return True
+        
     except Exception as e:
         print(f"❌ Model loading failed: {e}")
+        traceback.print_exc()
         return False
 
 def format_input(scenario_text):
@@ -69,11 +84,11 @@ def format_input(scenario_text):
         scenario = ", " + scenario.lstrip(", ")
     return f"Based on the situation, predict potential hazards and injuries. {scenario}<|endoftext|>"
 
-def generate_all_model_outputs(prompt, max_length=300, temperature=0.7):
-    outputs = []
-    for model, tokenizer in zip(models, tokenizers):
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-        inputs = {k: v.to(next(model.parameters()).device) for k, v in inputs.items()}
+def generate_single_model_output(model, tokenizer, prompt, max_length=300, temperature=0.7):
+    """Generate output from a single model"""
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+        inputs = {k: v.to('cpu') for k, v in inputs.items()}
 
         with torch.no_grad():
             output = model.generate(
@@ -84,207 +99,275 @@ def generate_all_model_outputs(prompt, max_length=300, temperature=0.7):
                 top_k=50,
                 repetition_penalty=1.1,
                 pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id
+                eos_token_id=tokenizer.eos_token_id,
+                do_sample=True
             )
 
         decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-        outputs.append(f"=== RAW RESPONSE START ===\n{decoded}\n=== RAW RESPONSE END ===")
-    return outputs
+        return decoded
+        
+    except Exception as e:
+        print(f"Error generating from model: {e}")
+        return f"Error: Unable to generate response - {str(e)}"
 
-def extract_scenario_from_prompt(prompt):
-    try:
-        return re.sub(r"^.*predict potential hazards and injuries\.\s*", "", prompt)
-    except:
-        return prompt
+def generate_all_model_outputs(prompt, max_length=300, temperature=0.7):
+    """Generate outputs from all loaded models"""
+    outputs = []
+    
+    for i, (model, tokenizer) in enumerate(zip(models, tokenizers)):
+        print(f"Generating from {model_names[i]}...")
+        
+        raw_output = generate_single_model_output(model, tokenizer, prompt, max_length, temperature)
+        formatted_output = f"=== {model_names[i].upper()} RAW RESPONSE START ===\n{raw_output}\n=== {model_names[i].upper()} RAW RESPONSE END ==="
+        outputs.append(formatted_output)
+        
+    return outputs
 
-def parse_json_from_raw_output(raw_output):
-    """Extract JSON from raw model output"""
+def analyze_with_cpu_model(raw_outputs, zero_shot_injury):
+    """Use CPU-only model for analysis instead of Gemini Pro"""
+    
+    # Initialize CPU analysis model (using a better model for reasoning)
     try:
-        # Look for JSON pattern in the raw output
-        json_match = re.search(r'\{.*?\}', raw_output, re.DOTALL)
-        if json_match:
-            json_str = json_match.group(0)
-            return json.loads(json_str)
-        return None
+        analysis_pipeline = pipeline(
+            "text-generation",
+            model="microsoft/phi-2",
+            device=-1,
+            torch_dtype=torch.float32,
+            trust_remote_code=True
+        )
     except:
-        return None
-
-def extract_structured_data_from_outputs(raw_outputs):
-    """Extract and combine structured JSON data from all model outputs"""
-    all_json_data = []
-    
-    for output in raw_outputs:
-        json_data = parse_json_from_raw_output(output)
-        if json_data:
-            all_json_data.append(json_data)
-    
-    return all_json_data
-
-def analyze_with_advanced_hf_model(raw_outputs, zero_shot_injury, structured_data):
-    """Replace Gemini Pro functionality with advanced HF model analysis"""
+        # Fallback to simpler model if Phi-2 fails
+        analysis_pipeline = pipeline(
+            "text-generation", 
+            model="gpt2",
+            device=-1
+        )
     
-    # Prepare the analysis prompt similar to original Gemini prompt
-    structured_info = ""
-    if structured_data:
-        structured_info = "\n\nStructured data extracted from models:\n"
-        for i, data in enumerate(structured_data, 1):
-            structured_info += f"Model {i}: {json.dumps(data, indent=2)}\n"
+    # Prepare concise prompt for CPU model
+    models_summary = ""
+    for i, output in enumerate(raw_outputs):
+        # Extract key information from each model output
+        clean_output = output.replace("=== RAW RESPONSE START ===", "").replace("=== RAW RESPONSE END ===", "").strip()
+        models_summary += f"Model {i+1}: {clean_output[:200]}...\n"
     
-    prompt = f"""You are a workplace safety analyst. Below are raw text outputs from construction safety prediction models.
-
-Your tasks:
-- Compare and merge the model outputs
-- Summarize the most plausible cause of accident in natural language
-- Infer the degree of injury by considering all outputs and classifier suggestion
-
-Classifier prediction for Degree of Injury: {zero_shot_injury}
+    prompt = f"""Analyze workplace safety incident based on model predictions:
 
-Model Outputs:
-{raw_outputs[0]}
+{models_summary}
 
-{raw_outputs[1] if len(raw_outputs) > 1 else ""}
+Classifier suggests: {zero_shot_injury} severity
 
-{raw_outputs[2] if len(raw_outputs) > 2 else ""}
-
-{structured_info}
-
-Based on this analysis, provide a concise response in this format:
-Cause of Accident: [single clear sentence]
-Degree of Injury: [Low/Medium/High]
-
-Analysis:"""
+Task: Integrate the model outputs to identify the main cause and injury level.
+Response format:
+Cause of Accident:"""
 
     try:
-        # Use the analysis model to generate response
-        response = analysis_model(
-            prompt, 
-            max_length=len(prompt.split()) + 100,
-            temperature=0.3,  # Lower temperature for more consistent analysis
+        # Generate analysis
+        response = analysis_pipeline(
+            prompt,
+            max_length=len(prompt.split()) + 80,
+            temperature=0.3,
             do_sample=True,
-            pad_token_id=analysis_model.tokenizer.eos_token_id
+            pad_token_id=analysis_pipeline.tokenizer.eos_token_id,
+            truncation=True
         )
         
         generated_text = response[0]['generated_text']
-        # Extract only the generated part after the prompt
         analysis_result = generated_text.replace(prompt, "").strip()
         
-        # If the analysis doesn't contain the required format, create it
-        if "Cause of Accident:" not in analysis_result:
-            # Fallback analysis based on structured data
-            cause = "Multiple safety protocol violations identified"
-            if structured_data:
-                causes = []
-                for data in structured_data:
-                    if isinstance(data, dict) and "Cause of Accident" in data:
-                        causes.append(data["Cause of Accident"])
-                if causes:
-                    cause = causes[0]  # Take the first cause found
-            
-            analysis_result = f"Cause of Accident: {cause}\nDegree of Injury: {zero_shot_injury}"
+        # Ensure proper format
+        if not analysis_result or len(analysis_result) < 10:
+            return perform_rule_based_analysis(raw_outputs, zero_shot_injury)
         
-        return analysis_result
+        # Add degree of injury if not present
+        if "Degree of Injury:" not in analysis_result:
+            analysis_result += f"\nDegree of Injury: {zero_shot_injury}"
+        
+        return f"Cause of Accident: {analysis_result}\nDegree of Injury: {zero_shot_injury}"
         
     except Exception as e:
-        print("❌ Advanced HF model analysis failed:", e)
-        # Fallback using structured data if available
-        if structured_data and len(structured_data) > 0:
-            first_data = structured_data[0]
-            cause = first_data.get("Cause of Accident", "Safety protocol violation")
-            injury = first_data.get("Degree of Injury", zero_shot_injury)
-            return f"Cause of Accident: {cause}\nDegree of Injury: {injury}"
+        print(f"❌ CPU model analysis failed: {e}")
+        return perform_rule_based_analysis(raw_outputs, zero_shot_injury)
+
+def perform_rule_based_analysis(raw_outputs, zero_shot_injury):
+    """Fallback rule-based analysis when CPU model fails"""
+    
+    # Combine all model outputs
+    all_text = " ".join(raw_outputs).lower()
+    
+    # Define safety categories and their indicators
+    safety_categories = {
+        'fall_protection': ['fall', 'height', 'scaffolding', 'ladder', 'roof', 'protection'],
+        'chemical_exposure': ['chemical', 'solvent', 'toxic', 'fumes', 'vapor', 'exposure'],
+        'equipment_failure': ['equipment', 'machinery', 'malfunction', 'failure', 'maintenance'],
+        'fire_safety': ['fire', 'ignition', 'flammable', 'welding', 'spark', 'combustible'],
+        'electrical': ['electrical', 'shock', 'current', 'wire', 'power'],
+        'confined_space': ['confined', 'space', 'ventilation', 'oxygen', 'gas']
+    }
+    
+    # Score each category
+    category_scores = {}
+    for category, keywords in safety_categories.items():
+        score = sum(1 for keyword in keywords if keyword in all_text)
+        if score > 0:
+            category_scores[category] = score
+    
+    # Determine primary cause
+    if category_scores:
+        primary_category = max(category_scores, key=category_scores.get)
+        
+        cause_descriptions = {
+            'fall_protection': "Inadequate fall protection measures leading to worker falling from height",
+            'chemical_exposure': "Unsafe chemical handling without proper protective equipment causing exposure",
+            'equipment_failure': "Equipment malfunction due to inadequate maintenance or safety protocols",
+            'fire_safety': "Fire safety protocol violations resulting in ignition of flammable materials",
+            'electrical': "Electrical safety hazards due to improper handling or faulty equipment",
+            'confined_space': "Confined space entry without proper safety procedures and ventilation"
+        }
+        
+        primary_cause = cause_descriptions.get(primary_category, "Multiple safety protocol violations")
+        
+        # Add secondary factors if present
+        secondary_factors = [cat for cat, score in category_scores.items() 
+                           if cat != primary_category and score > 0]
         
-        return f"Cause of Accident: Unable to analyze due to technical error\nDegree of Injury: {zero_shot_injury}"
+        if secondary_factors:
+            primary_cause += f". Contributing factors include {', '.join(secondary_factors[:2])} safety issues"
+        
+    else:
+        primary_cause = "Safety incident due to inadequate risk assessment and protocol violations"
+    
+    return f"Cause of Accident: {primary_cause}.\nDegree of Injury: {zero_shot_injury}"
+
+def extract_scenario_from_prompt(prompt):
+    try:
+        return re.sub(r"^.*predict potential hazards and injuries\.\s*", "", prompt)
+    except:
+        return prompt
 
 def generate_prediction_ensemble(scenario_text, max_length=300, temperature=0.7):
+    """Main prediction function using CPU-only models"""
+    
     if not scenario_text.strip():
-        return "❌ Please enter a scenario.", "", "", ""
+        return "❌ Please enter a scenario.", "", ""
 
     try:
+        # Generate prompt
         prompt = format_input(scenario_text)
+        
+        # Generate outputs from all models
+        print("Generating outputs from all models...")
         raw_outputs = generate_all_model_outputs(prompt, max_length, temperature)
-
+        
+        # Get zero-shot classification
         scenario_only = extract_scenario_from_prompt(prompt)
         injury_guess = classify_injury_zero_shot(scenario_only)
-
-        # Extract structured JSON data from raw outputs
-        structured_data = extract_structured_data_from_outputs(raw_outputs)
-
-        # Use advanced HF model analysis (replacing Gemini)
-        hf_analysis = analyze_with_advanced_hf_model(raw_outputs, injury_guess, structured_data)
-
-        # Parse the analysis results
-        match_cause = re.search(r"Cause of Accident\s*:\s*(.+)", hf_analysis)
-        match_injury = re.search(r"Degree of Injury\s*:\s*(Low|Medium|High)", hf_analysis, re.IGNORECASE)
-
-        cause = match_cause.group(1).strip() if match_cause else "Unable to determine cause"
+        
+        # Use CPU model for analysis instead of Gemini
+        print("Analyzing with CPU model...")
+        cpu_analysis = analyze_with_cpu_model(raw_outputs, injury_guess)
+        
+        # Parse CPU analysis response
+        match_cause = re.search(r"Cause of Accident\s*:\s*(.+?)(?=\nDegree of Injury|$)", cpu_analysis, re.DOTALL)
+        match_injury = re.search(r"Degree of Injury\s*:\s*(Low|Medium|High)", cpu_analysis, re.IGNORECASE)
+        
+        cause = match_cause.group(1).strip() if match_cause else "Unable to determine cause from model outputs"
         injury = match_injury.group(1).strip().capitalize() if match_injury else injury_guess
-
+        
         combined_raw = "\n\n".join(raw_outputs)
         
-        # Format structured data for display
-        structured_display = json.dumps(structured_data, indent=2) if structured_data else "No structured data found"
+        return cause, injury, combined_raw
         
-        return cause, injury, combined_raw, structured_display
-
     except Exception as e:
-        return "❌ Prediction failed.", "", traceback.format_exc(), ""
+        error_msg = f"❌ Prediction failed: {str(e)}"
+        print(error_msg)
+        traceback.print_exc()
+        return error_msg, "", ""
 
 def create_interface():
-    with gr.Blocks(title="Workplace Safety Risk Predictor") as interface:
-        gr.HTML("""
-        <h1>🚧 Workplace Safety Risk Prediction Model (Enhanced Ensemble)</h1>
-        <p>Enter a construction scenario to analyze possible risks. Uses advanced language models for better analysis.</p>
-        <p><strong>Expected JSON Output Format:</strong></p>
-        <pre>{"Cause of Accident": "...", "Degree of Injury": "High/Medium/Low", "Hazards": ["...", "..."]}</pre>
-        <p><strong>Examples:</strong></p>
+    with gr.Blocks(title="Multi-Model Safety Risk Predictor") as interface:
+        gr.HTML(f"""
+        <h1>🚧 Multi-Model Safety Risk Predictor (CPU-Only)</h1>
+        <p><strong>System Overview:</strong></p>
         <ul>
-            <li>An employee was working with chemical solvents without proper ventilation. The employee inhaled toxic fumes and experienced respiratory problems.</li>
-            <li>A worker fell from scaffolding due to lack of fall protection measures in place.</li>
-            <li>While operating a crane, the load became unstable and struck a nearby worker.</li>
-            <li>During welding, flammable vapors ignited due to poor fire safety practices.</li>
+            <li>Loads {len(MODEL_PATHS)} specialized safety prediction models</li>
+            <li>Each model analyzes the scenario independently</li>
+            <li>CPU-only analysis model integrates all results using advanced reasoning</li>
+            <li>Handles conflicting predictions through pattern analysis and majority consensus</li>
+            <li>Fully optimized for CPU-only Hugging Face Spaces</li>
         </ul>
+        <p><strong>Models Loaded:</strong> {len(models)} / {len(MODEL_PATHS)}</p>
+        <p><strong>Base Model:</strong> {BASE_MODEL_ID}</p>
+        <p><strong>Analysis Method:</strong> CPU-Only (No external API calls)</p>
         """)
 
         with gr.Row():
             with gr.Column():
-                scenario_input = gr.Textbox(lines=5, label="Scenario Description")
+                scenario_input = gr.Textbox(
+                    lines=6, 
+                    label="Construction Scenario Description",
+                    placeholder="Describe the workplace safety incident or scenario..."
+                )
+                
                 gr.Markdown("**Quick Examples:**")
                 with gr.Row():
-                    ex1 = gr.Button("Solvent Exposure")
-                    ex2 = gr.Button("Fall from Scaffolding")
-                    ex3 = gr.Button("Crane Load Accident")
-                    ex4 = gr.Button("Welding Fire Hazard")
-                temperature = gr.Slider(0.1, 1.0, 0.7, 0.1, label="Creativity (Temperature)")
-                max_len = gr.Slider(100, 500, 300, 50, label="Max Response Length")
-                predict_btn = gr.Button("🔍 Analyze")
+                    ex1 = gr.Button("Chemical Exposure", size="sm")
+                    ex2 = gr.Button("Fall Hazard", size="sm") 
+                    ex3 = gr.Button("Equipment Malfunction", size="sm")
+                    ex4 = gr.Button("Fire Incident", size="sm")
+                
+                with gr.Row():
+                    temperature = gr.Slider(0.1, 1.0, 0.7, 0.1, label="Model Creativity")
+                    max_len = gr.Slider(100, 400, 300, 50, label="Response Length")
+                
+                predict_btn = gr.Button("🔍 Analyze with Multi-Model Ensemble", variant="primary")
 
             with gr.Column():
-                cause_output = gr.Textbox(label="📝 Cause of Accident")
-                degree_output = gr.Textbox(label="📈 Degree of Injury")
-                with gr.Accordion("📊 Extracted Structured Data", open=False):
-                    structured_output = gr.Textbox(label="JSON Data from Models", lines=8)
-                with gr.Accordion("📄 Raw Model Outputs", open=False):
-                    raw_output = gr.Textbox(label="Raw Responses", lines=12)
-
+                cause_output = gr.Textbox(
+                    label="📝 Integrated Cause Analysis", 
+                    lines=4,
+                    info="CPU model's integrated analysis of all model outputs"
+                )
+                degree_output = gr.Textbox(
+                    label="📈 Degree of Injury",
+                    info="Based on zero-shot classification + model integration"
+                )
+                
+                with gr.Accordion("📄 Individual Model Outputs", open=False):
+                    raw_output = gr.Textbox(label="Raw Model Responses", lines=15)
+
+        # Event handlers
         predict_btn.click(
             fn=generate_prediction_ensemble,
             inputs=[scenario_input, max_len, temperature],
-            outputs=[cause_output, degree_output, raw_output, structured_output]
+            outputs=[cause_output, degree_output, raw_output]
         )
 
-        ex1.click(fn=lambda: "An employee was working with chemical solvents without proper ventilation. The employee inhaled toxic fumes and experienced respiratory problems.", outputs=scenario_input)
-        ex2.click(fn=lambda: "A worker fell from scaffolding due to lack of fall protection measures in place.", outputs=scenario_input)
-        ex3.click(fn=lambda: "While operating a crane, the load became unstable and struck a nearby worker.", outputs=scenario_input)
-        ex4.click(fn=lambda: "During welding, flammable vapors ignited due to poor fire safety practices.", outputs=scenario_input)
+        # Example scenarios
+        ex1.click(fn=lambda: "An employee was working with chemical solvents in a poorly ventilated area without proper respiratory protection. The worker began experiencing dizziness and respiratory distress after prolonged exposure.", outputs=scenario_input)
+        
+        ex2.click(fn=lambda: "A construction worker was installing roofing materials on a steep slope without proper fall protection equipment. The worker lost footing on wet materials and fell.", outputs=scenario_input)
+        
+        ex3.click(fn=lambda: "During routine maintenance, a hydraulic press malfunctioned due to worn seals. The operator's hand was caught when the press unexpectedly activated.", outputs=scenario_input)
+        
+        ex4.click(fn=lambda: "While welding in an area with flammable materials, proper fire safety protocols were not followed. Sparks ignited nearby combustible materials causing a flash fire.", outputs=scenario_input)
 
-        gr.HTML("<p style='text-align:center;'>Built with Advanced Transformers + Enhanced Analysis + Gradio</p>")
+        gr.HTML(f"""
+        <div style='text-align:center; margin-top:20px;'>
+            <p><strong>System Status:</strong> {len(models)} models loaded | CPU-optimized | No external APIs</p>
+            <p><em>Built with Multi-Model Ensemble + CPU Analysis + Gradio</em></p>
+        </div>
+        """)
 
     return interface
 
-print("🚀 Starting app...")
+# Initialize and launch
+print("🚀 Starting Multi-Model Safety Predictor...")
+print(f"Attempting to load {len(MODEL_PATHS)} models...")
+
 if load_models():
+    print(f"✅ Successfully loaded {len(models)} models")
     app = create_interface()
     if __name__ == "__main__":
         app.launch(server_name="0.0.0.0", server_port=7860, share=True)
 else:
-    print("❌ Failed to load models.")
+    print("❌ Failed to load models. Please check model paths and system resources.")