Update app.py

app.py

@@ -1,6 +1,9 @@
 import torch
 import re
 import gradio as gr
+import json
+import ast
+import traceback
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
@@ -52,31 +55,28 @@ def classify_injury_zero_shot(description):
     result = injury_classifier(description, candidate_labels)
     return label_map[result["labels"][0]]
 
+# === FORMAT INPUT ===
+def format_input(scenario_text):
+    scenario = scenario_text.strip()
+    if not scenario.startswith(", "):
+        scenario = scenario.lstrip(", ")
+    return f"Based on the situation, predict potential hazards and injuries. {scenario}<|endoftext|>"
+
 # === GENERATION FROM EACH MODEL ===
 def generate_single_model_output(model, tokenizer, prompt, max_length=300, temperature=0.7):
-    try:
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to("cpu")
-
-        with torch.no_grad():
-            output = model.generate(
-                **inputs,
-                max_length=inputs["input_ids"].shape[1] + max_length,
-                temperature=temperature,
-                top_p=0.9,
-                top_k=50,
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id,
-                do_sample=True
-            )
-
-        if output is not None and len(output) > 0:
-            decoded = tokenizer.decode(output[0], skip_special_tokens=True).strip()
-            return decoded
-        else:
-            return "[No output was generated by the model.]"
-
-    except Exception as e:
-        return f"[Error generating output: {str(e)}]"
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to("cpu")
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_length=inputs["input_ids"].shape[1] + max_length,
+            temperature=temperature,
+            top_p=0.9,
+            top_k=50,
+            repetition_penalty=1.1,
+            pad_token_id=tokenizer.eos_token_id,
+            do_sample=True
+        )
+    return tokenizer.decode(output[0], skip_special_tokens=True).strip()
 
 # === ANALYSIS WITH FLAN-T5 ===
 def analyze_with_cpu_model(raw_outputs, zero_shot_injury):
@@ -114,7 +114,7 @@ def generate_prediction_ensemble(scenario_text, max_len, temperature):
     if not scenario_text.strip():
         return "Please enter a scenario", "", ""
 
-    prompt = f"Based on the situation, predict potential hazards and injuries.\n{scenario_text.strip()}"
+    prompt = format_input(scenario_text)
 
     raw_outputs = [
         generate_single_model_output(model, tokenizer, prompt, max_length=max_len, temperature=temperature)