Update app.py

app.py

@@ -1,83 +1,350 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import json
 import re
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+import warnings
+warnings.filterwarnings("ignore")
 
-# Models dictionary
-model_paths = {
-    "Model A (OSHA LLaMA)": "FrAnKu34t23/Construction_Mistral_Risk_Prediction_Model_v1",
-    "Model B (OSHA LLaMA2)": "FrAnKu34t23/Construction_Mistral_Risk_Prediction_Model_v2"
-}
-
-def build_prompt(desc):
-    return f"""Incident Description:
-{desc}
-
-Please extract the following in JSON format:
-- Hazards
-- Cause of Accident
-- Degree of Injury (High, Medium, Low)
-- Occupation
+# Configuration
+BASE_MODEL_ID = "distilgpt2"
+LORA_MODEL_ID = "FrAnKu34t23/Construction_Mistral_Risk_Prediction_Model_v3"
 
-Output must be valid JSON using double quotes.
-"""
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
 
-def parse_response(output):
+def load_model():
+    """Load the base model and LoRA adapter"""
+    global model, tokenizer
+    
     try:
-        match = re.search(r'\{[\s\S]*?\}', output)
-        if match:
-            data = json.loads(match.group().replace("'", '"'))
-            return json.dumps(data, indent=2)
-        else:
-            return "❌ Could not parse JSON from output."
+        print("🔍 Loading base model and tokenizer...")
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Load base model
+        base_model = AutoModelForCausalLM.from_pretrained(
+            BASE_MODEL_ID,
+            torch_dtype=torch.float32,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        
+        print("📁 Loading LoRA adapter...")
+        
+        # Load LoRA adapter from Hugging Face
+        model = PeftModel.from_pretrained(base_model, LORA_MODEL_ID)
+        model.eval()
+        
+        print("✅ Model loaded successfully!")
+        return True
+        
     except Exception as e:
-        return f"❌ Error parsing output: {e}"
-
-def generate(incident, model_choice):
-    print(incident, model_choice)
-    model_id = model_paths[model_choice]
-
-    # Load model/tokenizer on demand
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map=None,        # no device map
-        torch_dtype=torch.float32,  # force CPU-friendly dtype
-        low_cpu_mem_usage=True
-    ).to("cpu")
+        print(f"❌ Error loading model: {e}")
+        return False
 
-    model.eval()
+def format_input(scenario_text):
+    """Format input to match training data format"""
+    # Ensure the input starts with ", " like training data
+    cleaned_text = scenario_text.strip()
+    if not cleaned_text.startswith(", "):
+        if cleaned_text.startswith(","):
+            cleaned_text = ", " + cleaned_text[1:].strip()
+        else:
+            cleaned_text = ", " + cleaned_text
+    
+    instruction = "Based on the situation, predict potential hazards and injuries."
+    formatted_prompt = f"{instruction} {cleaned_text}"
+    
+    return formatted_prompt
 
-    prompt = build_prompt(incident)
-    inputs = tokenizer(prompt, return_tensors='pt').to("cpu")
+def parse_json_response(response_text):
+    """Extract and parse JSON from model response"""
+    try:
+        # First, try to parse the entire response as JSON
+        if response_text.strip().startswith('{') and response_text.strip().endswith('}'):
+            return json.loads(response_text.strip())
+        
+        # If that fails, look for JSON pattern in the text
+        json_pattern = r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}'
+        matches = re.findall(json_pattern, response_text, re.DOTALL)
+        
+        for match in matches:
+            try:
+                return json.loads(match)
+            except:
+                continue
+        
+        # If no valid JSON found, return structured error
+        return {
+            "Hazards": ["Unable to parse response"],
+            "Cause of Accident": "Model output parsing failed",
+            "Degree of Injury": "Unknown",
+            "raw_response": response_text
+        }
+        
+    except Exception as e:
+        return {
+            "Hazards": [f"Parsing error: {str(e)}"],
+            "Cause of Accident": "JSON parsing failed",
+            "Degree of Injury": "Unknown",
+            "raw_response": response_text
+        }
 
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.9,
-            do_sample=True
+def generate_prediction(scenario_text, max_length=300, temperature=0.7):
+    """Generate workplace safety prediction"""
+    global model, tokenizer
+    
+    if model is None or tokenizer is None:
+        return "❌ Model not loaded. Please wait for initialization.", "", "", "", ""
+    
+    try:
+        # Format the input
+        formatted_prompt = format_input(scenario_text)
+        full_prompt = f"{formatted_prompt}{tokenizer.eos_token}"
+        
+        # Tokenize
+        inputs = tokenizer(
+            full_prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512,
+            padding=False
         )
+        
+        # Move to same device as model
+        device = next(model.parameters()).device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=len(inputs['input_ids'][0]) + max_length,
+                temperature=temperature,
+                do_sample=True,
+                top_p=0.9,
+                top_k=50,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                num_return_sequences=1,
+                repetition_penalty=1.1,
+                early_stopping=True
+            )
+        
+        # Decode response
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+        
+        # Extract generated part
+        input_text = tokenizer.decode(inputs['input_ids'][0], skip_special_tokens=False)
+        
+        if full_response.startswith(input_text):
+            generated_part = full_response[len(input_text):].strip()
+        else:
+            generated_part = full_response.strip()
+        
+        # Clean up response
+        if generated_part.startswith(tokenizer.eos_token):
+            generated_part = generated_part[len(tokenizer.eos_token):].strip()
+        
+        if generated_part.endswith(tokenizer.eos_token):
+            generated_part = generated_part[:-len(tokenizer.eos_token)].strip()
+        
+        # Parse the JSON response
+        parsed_response = parse_json_response(generated_part)
+        
+        # Extract individual components
+        hazards = parsed_response.get("Hazards", [])
+        cause = parsed_response.get("Cause of Accident", "Not specified")
+        degree = parsed_response.get("Degree of Injury", "Not specified")
+        
+        # Format hazards for display
+        hazards_display = ", ".join(hazards) if isinstance(hazards, list) else str(hazards)
+        
+        # Create formatted output
+        formatted_output = json.dumps(parsed_response, indent=2, ensure_ascii=False)
+        
+        return hazards_display, cause, degree, formatted_output, generated_part
+        
+    except Exception as e:
+        error_msg = f"❌ Error generating prediction: {str(e)}"
+        return error_msg, "", "", "", ""
 
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    parsed = parse_response(decoded)
-    return parsed
+def create_interface():
+    """Create the Gradio interface"""
+    
+    # Custom CSS for better styling
+    css = """
+    .gradio-container {
+        font-family: 'Arial', sans-serif;
+    }
+    .header {
+        text-align: center;
+        margin-bottom: 30px;
+    }
+    .warning-box {
+        background-color: #fff3cd;
+        border: 1px solid #ffeaa7;
+        border-radius: 5px;
+        padding: 15px;
+        margin: 10px 0;
+    }
+    """
+    
+    with gr.Blocks(css=css, title="Workplace Safety Risk Predictor") as interface:
+        
+        gr.HTML("""
+        <div class="header">
+            <h1>🚧 Workplace Safety Risk Prediction Model</h1>
+            <p>Analyze workplace scenarios to identify potential hazards, causes, and injury severity</p>
+        </div>
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                gr.HTML("<h3>📝 Enter Workplace Scenario</h3>")
+                
+                scenario_input = gr.Textbox(
+                    lines=5,
+                    placeholder="Example: an employee was operating a 400 ton mechanical power press. The press was actuated while the employee's right hand was in the point of operation...",
+                    label="Workplace Incident Description",
+                    info="Describe the workplace scenario you want to analyze"
+                )
+                
+                with gr.Row():
+                    with gr.Column():
+                        temperature = gr.Slider(
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=0.7,
+                            step=0.1,
+                            label="Creativity (Temperature)",
+                            info="Higher values = more creative responses"
+                        )
+                    
+                    with gr.Column():
+                        max_length = gr.Slider(
+                            minimum=100,
+                            maximum=500,
+                            value=300,
+                            step=50,
+                            label="Max Response Length",
+                            info="Maximum length of generated response"
+                        )
+                
+                predict_btn = gr.Button("🔍 Analyze Scenario", variant="primary", size="lg")
+                
+                gr.HTML("""
+                <div class="warning-box">
+                    <strong>⚠️ Note:</strong> This is an AI model for educational purposes. 
+                    Always consult safety professionals for real workplace safety assessments.
+                </div>
+                """)
+            
+            with gr.Column(scale=2):
+                gr.HTML("<h3>📊 Analysis Results</h3>")
+                
+                with gr.Row():
+                    with gr.Column():
+                        hazards_output = gr.Textbox(
+                            label="🚨 Identified Hazards",
+                            info="Potential hazards identified in the scenario"
+                        )
+                        
+                        cause_output = gr.Textbox(
+                            label="🔍 Cause of Accident",
+                            info="Primary cause classification"
+                        )
+                        
+                        degree_output = gr.Textbox(
+                            label="📈 Degree of Injury",
+                            info="Severity assessment"
+                        )
+                
+                with gr.Accordion("📋 Detailed JSON Output", open=False):
+                    json_output = gr.Code(
+                        label="Structured Response",
+                        language="json"
+                    )
+                
+                with gr.Accordion("🔍 Raw Model Output", open=False):
+                    raw_output = gr.Textbox(
+                        label="Raw Response",
+                        lines=3,
+                        info="Unprocessed model output"
+                    )
+        
+        # Example scenarios
+        gr.HTML("<h3>💡 Example Scenarios</h3>")
+        
+        with gr.Row():
+            example1 = gr.Button("Power Press Accident")
+            example2 = gr.Button("Fall from Ladder")
+            example3 = gr.Button("Chemical Exposure")
+            example4 = gr.Button("Lifting Injury")
+        
+        # Event handlers
+        predict_btn.click(
+            fn=generate_prediction,
+            inputs=[scenario_input, max_length, temperature],
+            outputs=[hazards_output, cause_output, degree_output, json_output, raw_output]
+        )
+        
+        # Example scenarios
+        example1.click(
+            lambda: "an employee was operating a 400 ton mechanical power press. The press was actuated while the employee's right hand was in the point of operation. The employee's fingers were amputated.",
+            outputs=scenario_input
+        )
+        
+        example2.click(
+            lambda: "an employee was using a ladder to access high shelves. The ladder was not properly secured and the employee fell from a height of 8 feet, resulting in head injuries.",
+            outputs=scenario_input
+        )
+        
+        example3.click(
+            lambda: "an employee was working with chemical solvents without proper ventilation. The employee inhaled toxic fumes and experienced respiratory problems.",
+            outputs=scenario_input
+        )
+        
+        example4.click(
+            lambda: "an employee was manually lifting heavy boxes weighing over 50 pounds without proper lifting technique or mechanical aids. The employee strained their back.",
+            outputs=scenario_input
+        )
+        
+        gr.HTML("""
+        <div style="text-align: center; margin-top: 30px; color: #666;">
+            <p>Built with ❤️ using Hugging Face Transformers and Gradio</p>
+            <p>Model: <a href="https://huggingface.co/FrAnKu34t23/Construction_Mistral_Risk_Prediction_Model_v3">Construction_Mistral_Risk_Prediction_Model_v3</a></p>
+        </div>
+        """)
+    
+    return interface
 
-demo = gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.Textbox(label="Incident Description", lines=6, placeholder="Describe the incident..."),
-        gr.Dropdown(choices=list(model_paths.keys()), 
-                    label="Choose Model",
-                    value=list(model_paths.keys())[0]  # ✅ Set default value
-                   )
-    ],
-    outputs=gr.Textbox(label="Extracted JSON"),
-    title="OSHA Risk Analyzer (CPU)",
-    description="Runs one OSHA model on CPU and extracts hazards, cause, injury level, and occupation."
-)
+# Initialize the model when the app starts
+print("🚀 Initializing Workplace Safety Risk Prediction App...")
+model_loaded = load_model()
 
-if __name__ == "__main__":
-    demo.launch()
+if model_loaded:
+    print("✅ App ready!")
+    # Create and launch the interface
+    app = create_interface()
+    
+    if __name__ == "__main__":
+        app.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=True
+        )
+else:
+    print("❌ Failed to load model. App cannot start.")
+    # Create a simple error interface
+    with gr.Blocks() as error_app:
+        gr.HTML("<h1>❌ Model Loading Failed</h1><p>Unable to load the safety prediction model.</p>")
+    
+    if __name__ == "__main__":
+        error_app.launch()