Update app.py

app.py

@@ -1,73 +1,124 @@
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import random
 import gradio as gr
 
-# Load the Hugging Face model and tokenizer
-model_name = "EleutherAI/gpt-neo-1.3B"  # Replace with your preferred model
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)  # Use CPU for Hugging Face Spaces
+# Model options
+model_options = {
+    "distilgpt2": "distilgpt2",
+    "GPT-Neo 125M": "EleutherAI/gpt-neo-125M",
+}
+
+# Load default model
+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
+
+# Predefined options for randomization
+names = ["John Doe", "Jane Smith", "Ali Khan"]
+locations = ["Pump House 1", "Main Valve Station", "Chemical Storage Area"]
+work_types = ["Routine pump maintenance", "Valve inspection", "Chemical handling"]
+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():
-    prompt = """
-    Write a detailed Safety Observation and Conversation (SOC) report for a water injection plant. The report should include:
-
-    1. Key Safety Conclusions/Comments/Agreements Made:
-    A brief summary of the observations, highlighting the main safety concerns and any corrective actions taken.
-
-    2. Observer's Name:
-    Provide a placeholder for the observer's name.
-
-    3. KOC ID No.:
-    Include a placeholder for the observer's ID.
+def generate_soc(model_choice, seed=None):
+    # Set seed if provided
+    if seed:
+        set_seed(seed)
+
+    # Update the generator if model_choice changes
+    global generator
+    model_name = model_options[model_choice]
+    if generator.tokenizer.name_or_path != model_name:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+
+    # Randomized fields
+    observer_name = random.choice(names)
+    location = random.choice(locations)
+    work_type = random.choice(work_types)
+    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)))
+    )
 
-    4. Type of Work Observed:
-    Specify the activity being observed (e.g., pump maintenance, valve inspection, or chemical handling).
+    # AI Prompt
+    prompt = f"""
+    Write a detailed Safety Observation and Conversation (SOC) report for a water injection plant. 
 
-    5. Location:
-    Mention the location within the plant where the observation occurred.
+    Key Safety Conclusions/Comments/Agreements Made:
+    Briefly summarize safety observations, key concerns, and corrective actions.
 
-    6. Duration (in mins):
-    Provide a placeholder for the time spent observing.
+    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:
-    Include detailed observations such as energy sources control, housekeeping, and spill containment. Each observation should be marked as Good Practice or Deviation with explanations for deviations.
+    {observations}
 
     --- People Observations:
-    Include details of personnel behavior, PPE compliance, hazard understanding, and examples of Good Practices or Deviations with corrective actions.
+    Include details on PPE compliance, hazard understanding, and good practices or deviations.
 
     --- Process Observations:
-    Include observations on job safety analysis, procedures followed, control of work standards, and improvements needed.
+    Summarize job safety analysis, procedures followed, and improvements needed.
 
     --- Performance Observations:
-    Evaluate the overall safety performance, including work pace, supervision, and prioritization of safety.
-
-    Format the output neatly with headers for each section and placeholders where needed.
+    Evaluate the overall safety performance, including work pace and supervision.
     """
-    result = generator(prompt, max_length=1024, num_return_sequences=1)[0]["generated_text"]
+    result = generator(prompt, max_length=512, num_return_sequences=1)[0]["generated_text"]
     return result
 
 # Gradio Interface
-def app_interface():
-    return generate_soc()
+def app_interface(model_choice, seed):
+    return generate_soc(model_choice, seed)
 
 # Gradio Layout
 with gr.Blocks() as app:
     gr.Markdown("# AI-Generated Safety Observation and Conversation (SOC) Reports")
     gr.Markdown(
         """
-        This app generates detailed SOC reports for a water injection plant. AI assists in creating observations, including conclusions, corrective actions, and key performance evaluations.
+        Generate detailed SOC reports for a water injection plant using AI assistance. 
+        Customize your report with multiple models, randomization, and reproducibility through seeds.
         """
     )
+
+    with gr.Row():
+        model_choice = gr.Dropdown(
+            label="Select Model",
+            choices=list(model_options.keys()),
+            value="GPT-Neo 125M",
+        )
+        seed = gr.Number(label="Seed (Optional)", value=None, precision=0)
+
     output_box = gr.Textbox(
         label="Generated SOC Report",
         placeholder="Your SOC report will appear here...",
         lines=30,
     )
-    generate_btn = gr.Button("Generate SOC Report")
-    copy_btn = gr.Button("Copy to Clipboard")
 
-    generate_btn.click(app_interface, outputs=output_box)
+    with gr.Row():
+        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)
     copy_btn.click(lambda text: text, inputs=output_box, outputs=None)
 
 # Launch the app