no message

main.py

@@ -140,7 +140,6 @@ def segment_text(text: str, max_tokens=500):  # Setting a conservative limit bel
 
 
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
-classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
 
 def robust_segment_text(text: str, max_tokens=510):  # Slightly less to ensure a buffer
     doc = nlp(text)
@@ -167,15 +166,18 @@ def robust_segment_text(text: str, max_tokens=510):  # Slightly less to ensure a
     return segments
 
 
+# Load a zero-shot classification model
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+
 def classify_segments(segments):
-    results = []
+    labels = ["Coverage Details", "Exclusions", "Premiums", "Claims Process", 
+              "Policy Limits", "Legal and Regulatory Information", "Renewals and Cancellations", 
+              "Discounts and Incentives", "Duties and Responsibilities", "Contact Information"]
+    classified_segments = []
     for segment in segments:
-        try:
-            result = classifier(segment)
-            results.append(result)
-        except Exception as e:
-            results.append({"error": str(e), "segment": segment[:50]})  # Include a part of the segment to debug if needed
-    return results
+        result = classifier(segment, candidate_labels=labels, multi_label=True)
+        classified_segments.append(result)
+    return classified_segments