Update app.py

app.py

@@ -13,19 +13,36 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 def format_output(text):
     impact, business, action = "N/A", "N/A", "N/A"
 
-    try:
-        lines = text.split("\n")
+    text = text.lower()
 
-        for line in lines:
-            if "Impact Level" in line:
-                impact = line.split(":")[-1].strip()
-            elif "Business Translation" in line:
-                business = line.split(":", 1)[-1].strip()
-            elif "Strategic Action" in line:
-                action = line.split(":", 1)[-1].strip()
+    # Try to extract sections flexibly
+    if "impact" in text:
+        try:
+            impact = text.split("impact")[1].split("\n")[0].replace("level", "").replace(":", "").strip().title()
+        except:
+            pass
 
-    except:
-        pass
+    if "business" in text:
+        try:
+            business = text.split("business")[1].split("strategic")[0].replace("translation", "").replace(":", "").strip().capitalize()
+        except:
+            pass
+
+    if "strategic" in text or "action" in text:
+        try:
+            action = text.split("strategic")[-1].replace("action", "").replace(":", "").strip().capitalize()
+        except:
+            pass
+
+    # FALLBACK: if parsing failed → use whole text smartly
+    if business == "N/A":
+        business = text[:120].capitalize()
+
+    if action == "N/A":
+        action = "Investigate the issue and apply appropriate fixes."
+
+    if impact == "N/A":
+        impact = "Moderate"
 
     return impact, business, action
     
@@ -51,12 +68,13 @@ Log:
 
     with torch.no_grad():
         outputs = model.generate(
-            **inputs,
-            max_new_tokens=120,
-            temperature=0.3,
-            top_p=0.9,
-            do_sample=True
-        )
+    **inputs,
+    max_new_tokens=100,
+    temperature=0.2,   # more stable
+    top_p=0.9,
+    do_sample=True,
+    repetition_penalty=1.3
+)
 
     result = tokenizer.decode(
     outputs[0][inputs["input_ids"].shape[-1]:],