Update app.py

app.py

@@ -1,11 +1,12 @@
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-import random
 import gradio as gr
+import random
 
-# Model options
+# Predefined model options
 model_options = {
     "distilgpt2": "distilgpt2",
     "GPT-Neo 125M": "EleutherAI/gpt-neo-125M",
+    "GPT-Neo 1.3B": "EleutherAI/gpt-neo-1.3B",
 }
 
 # Load default model
@@ -14,31 +15,22 @@ 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
 
-# Predefined options for randomization
-names = ["John Doe", "Jane Smith", "Ali Khan"]
-locations = ["Pump House 1", "Main Valve Station", "Chemical Storage Area"]
+# Random options for observations
+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"]
 durations = [30, 45, 60]
-good_practices = ["Good Practice"]
-deviations = ["Deviation"]
-
-plant_observations = [
-    ("Energy sources controlled", "Good Practice", "Lockout/tagout procedures were followed."),
-    ("Leaks/spills contained", "Deviation", "Oil spill near a pump flagged for cleanup."),
-    ("Housekeeping standard high", "Deviation", "Scattered tools were organized after reminder."),
-]
 
 # Function to set seed
 def set_seed(seed_value):
     random.seed(seed_value)
 
 # AI-based SOC report generation
-def generate_soc(model_choice, seed=None):
-    # Set seed if provided
+def generate_soc(model_choice, severity, seed=None):
     if seed:
         set_seed(seed)
 
-    # Update the generator if model_choice changes
+    # Update generator if model choice changes
     global generator
     model_name = model_options[model_choice]
     if generator.tokenizer.name_or_path != model_name:
@@ -46,49 +38,49 @@ def generate_soc(model_choice, seed=None):
         model = AutoModelForCausalLM.from_pretrained(model_name)
         generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
-    # Randomized fields
+    # Random selections for the fields
     observer_name = random.choice(names)
-    location = random.choice(locations)
     work_type = random.choice(work_types)
+    location = random.choice(locations)
     duration = random.choice(durations)
 
-    # Generate random plant observations
-    observations = "\n".join(
-        f"{i+1}. {obs[0]}\n{obs[1]}\n{obs[2]}"
-        for i, obs in enumerate(random.sample(plant_observations, len(plant_observations)))
-    )
+    # 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_text = severity_description.get(severity, "moderate concerns requiring attention.")
 
-    # AI Prompt
+    # AI prompt
     prompt = f"""
-    Write a detailed Safety Observation and Conversation (SOC) report for a water injection plant. 
-
-    Key Safety Conclusions/Comments/Agreements Made:
-    Briefly summarize safety observations, key concerns, and corrective actions.
+    Write a detailed Safety Observation and Conversation (SOC) report for a water injection plant with the following details:
 
     Observer's Name: {observer_name}
     KOC ID No.: [Insert KOC ID here]
     Type of Work Observed: {work_type}
     Location: {location}
-    Duration (in mins): {duration}
-
-    --- Plant Observations:
-    {observations}
+    Duration: {duration} minutes
 
-    --- People Observations:
-    Include details on PPE compliance, hazard understanding, and good practices or deviations.
+    Severity Level: {severity_text}
 
-    --- Process Observations:
-    Summarize job safety analysis, procedures followed, and improvements needed.
+    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).
 
-    --- Performance Observations:
-    Evaluate the overall safety performance, including work pace and supervision.
+    Format the output neatly in sections, and ensure it is professional and actionable.
     """
-    result = generator(prompt, max_length=512, num_return_sequences=1)[0]["generated_text"]
+
+    # Generate report using the selected model
+    result = generator(prompt, max_length=1024, num_return_sequences=1)[0]["generated_text"]
     return result
 
 # Gradio Interface
-def app_interface(model_choice, seed):
-    return generate_soc(model_choice, seed)
+def app_interface(model_choice, severity, seed=None):
+    return generate_soc(model_choice, severity, seed)
 
 # Gradio Layout
 with gr.Blocks() as app:
@@ -96,17 +88,24 @@ with gr.Blocks() as app:
     gr.Markdown(
         """
         Generate detailed SOC reports for a water injection plant using AI assistance. 
-        Customize your report with multiple models, randomization, and reproducibility through seeds.
+        Customize your report with multiple AI models, severity levels, and reproducibility using seeds.
         """
     )
 
     with gr.Row():
         model_choice = gr.Dropdown(
-            label="Select Model",
+            label="Select AI Model",
             choices=list(model_options.keys()),
             value="GPT-Neo 125M",
         )
-        seed = gr.Number(label="Seed (Optional)", value=None, precision=0)
+        severity_slider = gr.Slider(
+            label="Severity of SOC Report",
+            minimum=1,
+            maximum=3,
+            step=1,
+            value=2,
+        )
+        seed_input = gr.Number(label="Seed (Optional)", value=None, precision=0)
 
     output_box = gr.Textbox(
         label="Generated SOC Report",
@@ -118,7 +117,7 @@ with gr.Blocks() as app:
         generate_btn = gr.Button("Generate SOC Report")
         copy_btn = gr.Button("Copy to Clipboard")
 
-    generate_btn.click(app_interface, inputs=[model_choice, seed], outputs=output_box)
+    generate_btn.click(app_interface, inputs=[model_choice, severity_slider, seed_input], outputs=output_box)
     copy_btn.click(lambda text: text, inputs=output_box, outputs=None)
 
 # Launch the app