Update app.py

app.py

@@ -54,19 +54,29 @@ def classify_injury_zero_shot(description):
 
 # === GENERATION FROM EACH MODEL ===
 def generate_single_model_output(model, tokenizer, prompt, max_length=300, temperature=0.7):
-    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()
+    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)}]"
 
 # === ANALYSIS WITH FLAN-T5 ===
 def analyze_with_cpu_model(raw_outputs, zero_shot_injury):