2

app.py

@@ -4,25 +4,38 @@ from transformers import pipeline
 classifier = pipeline(
     "text-generation",
     model="AtrriJi/smolified-risk-clause-classifier",
-    device=-1  # force CPU
+    device=-1
 )
 
 def classify(clause):
+
     prompt = f"""
-Classify the legal clause into a category and risk level.
-Return JSON with keys category and risk_level.
+Classify the legal clause into one of:
+Payment Terms, Intellectual Property, Confidentiality,
+Termination, Indemnification, Force Majeure,
+Governing Law, Warranty, Limitation of Liability,
+Dispute Resolution.
+
+Also predict risk level: Low, Medium, or High.
+
+Return ONLY JSON like:
+{{"category": "...", "risk_level": "..."}}
 
 Clause:
 {clause}
 """
 
-    output = classifier(
+    result = classifier(
         prompt,
         max_new_tokens=100,
         do_sample=False
     )[0]["generated_text"]
 
-    return output
+    # Remove prompt part
+    clean_output = result[len(prompt):].strip()
+
+    return clean_output
+
 
 demo = gr.Interface(
     fn=classify,