Update app.py

app.py

@@ -2,7 +2,7 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 import random
 
-# Predefined model options
+# Model options
 model_options = {
     "distilgpt2": "distilgpt2",
     "GPT-Neo 125M": "EleutherAI/gpt-neo-125M",
@@ -13,24 +13,82 @@ model_options = {
 default_model_name = model_options["GPT-Neo 125M"]
 tokenizer = AutoTokenizer.from_pretrained(default_model_name)
 model = AutoModelForCausalLM.from_pretrained(default_model_name)
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)  # Use CPU
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
-# Random options for observations
+# Randomized options for report fields
 names = ["WPMPOperator 1", "John Doe", "Ali Khan"]
 work_types = ["Routine pump maintenance", "Valve inspection", "Chemical handling"]
-locations = ["Water Injection Plant - Pump House 2", "Main Valve Station", "Chemical Storage Area"]
+locations = ["Pump House 1", "Main Valve Station", "Chemical Storage Area"]
 durations = [30, 45, 60]
 
-# Function to set seed
+observations = {
+    "Plant": [
+        "Energy sources controlled",
+        "Plant well maintained",
+        "Leaks/spills contained",
+        "Protection from hazards in place",
+        "Access/egress clear",
+        "Layout and work locations safe",
+        "Housekeeping standard high",
+    ],
+    "People": [
+        "Work location protected",
+        "Work positions safe",
+        "People competent",
+        "Hazards understood",
+        "PPE appropriate",
+        "Risk to others avoided",
+        "Distractions absent",
+    ],
+    "Process": [
+        "Procedures valid",
+        "Correct procedures used",
+        "Control of work standard applied",
+        "Control of work requirements understood",
+        "Risk documented",
+        "Job Safety Analysis participation",
+        "Work well organized, systematic",
+        "Communication effective",
+        "Change managed safely",
+    ],
+    "Performance": [
+        "Work pressure does not compromise safety",
+        "Safety priority messaged by leadership",
+        "Supervision appropriate",
+        "Pace appropriate/safe",
+        "Safety performance recognized",
+    ],
+}
+
+# Random comments
+comments_pool = [
+    "No major issues observed.",
+    "Corrective actions were taken immediately.",
+    "Workers followed safety protocols.",
+    "Area was well-maintained.",
+    "Minor deviations flagged for correction.",
+    "Immediate action required for identified hazards.",
+]
+
+# Set random seed
 def set_seed(seed_value):
     random.seed(seed_value)
 
+# Generate observations with statuses
+def generate_section_observations(section_name):
+    results = []
+    for obs in observations[section_name]:
+        status = random.choice(["Good Practice", "Deviation", "NA"])
+        comment = random.choice(comments_pool) if status != "NA" else ""
+        results.append((obs, status, comment))
+    return results
+
 # AI-based SOC report generation
 def generate_soc(model_choice, severity, seed=None):
     if seed:
         set_seed(seed)
 
-    # Update generator if model choice changes
+    # Update model if selection changes
     global generator
     model_name = model_options[model_choice]
     if generator.tokenizer.name_or_path != model_name:
@@ -38,45 +96,58 @@ def generate_soc(model_choice, severity, seed=None):
         model = AutoModelForCausalLM.from_pretrained(model_name)
         generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
-    # Random selections for the fields
+    # Random selections for fields
     observer_name = random.choice(names)
     work_type = random.choice(work_types)
     location = random.choice(locations)
     duration = random.choice(durations)
 
-    # Adjust tone based on severity slider
-    severity_description = {
-        1: "minor concerns and deviations were observed.",
-        2: "moderate concerns requiring immediate attention were identified.",
-        3: "serious safety concerns were flagged for urgent corrective action."
+    # Severity description
+    severity_levels = {
+        1: "minor concerns and deviations.",
+        2: "moderate safety concerns requiring attention.",
+        3: "serious safety issues needing urgent action.",
     }
-    severity_text = severity_description.get(severity, "moderate concerns requiring attention.")
+    severity_text = severity_levels.get(severity, "moderate safety concerns.")
 
-    # AI prompt
+    # AI Prompt
     prompt = f"""
-    Write a detailed Safety Observation and Conversation (SOC) report for a water injection plant with the following details:
-
+    Generate a Safety Observation and Conversation (SOC) report for the following details:
     Observer's Name: {observer_name}
     KOC ID No.: [Insert KOC ID here]
     Type of Work Observed: {work_type}
     Location: {location}
     Duration: {duration} minutes
+    Severity: {severity_text}
+    Include:
+    - Key safety conclusions.
+    - Randomized observations for Plant, People, Process, and Performance sections.
+    Format output similar to the MyHSSE layout with numbered observations and status.
+    """
 
-    Severity Level: {severity_text}
+    # Generate the main report body
+    report = generator(prompt, max_length=1024, num_return_sequences=1)[0]["generated_text"]
 
-    The report should include:
-    - Key safety conclusions, concerns, and corrective actions taken.
-    - Plant observations (e.g., energy control, housekeeping) marked as Good Practice or Deviation with comments.
-    - People observations (e.g., PPE compliance, hazard understanding).
-    - Process observations (e.g., job safety analysis, procedures).
-    - Performance observations (e.g., pace, supervision, and safety prioritization).
+    # Random observations
+    plant_obs = generate_section_observations("Plant")
+    people_obs = generate_section_observations("People")
+    process_obs = generate_section_observations("Process")
+    performance_obs = generate_section_observations("Performance")
 
-    Format the output neatly in sections, and ensure it is professional and actionable.
-    """
+    def format_section(section_title, observations):
+        formatted = f"\n{section_title} Observations:\n"
+        formatted += "\n".join(
+            f"{i+1}. {obs[0]} - {obs[1]}\n   Comments: {obs[2]}" for i, obs in enumerate(observations)
+        )
+        return formatted
+
+    # Append formatted sections
+    report += format_section("Plant", plant_obs)
+    report += format_section("People", people_obs)
+    report += format_section("Process", process_obs)
+    report += format_section("Performance", performance_obs)
 
-    # Generate report using the selected model
-    result = generator(prompt, max_length=1024, num_return_sequences=1)[0]["generated_text"]
-    return result
+    return report
 
 # Gradio Interface
 def app_interface(model_choice, severity, seed=None):
@@ -87,8 +158,8 @@ with gr.Blocks() as app:
     gr.Markdown("# AI-Generated Safety Observation and Conversation (SOC) Reports")
     gr.Markdown(
         """
-        Generate detailed SOC reports for a water injection plant using AI assistance. 
-        Customize your report with multiple AI models, severity levels, and reproducibility using seeds.
+        Generate detailed SOC reports in the MyHSSE format using AI assistance.
+        Customize with multiple AI models, severity levels, and reproducibility using seeds.
         """
     )