no message

main.py

@@ -139,31 +139,50 @@ def segment_text(text: str, max_tokens=500):  # Setting a conservative limit bel
     return segments
 
 
-classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
+tokenizer = AutoTokenizer.from_pretrained("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)
+    segments = []
+    current_segment = []
+    current_tokens = []
+
+    for sent in doc.sents:
+        words = sent.text.strip().split()
+        sentence_tokens = tokenizer.encode(' '.join(words), add_special_tokens=False)
+        
+        if len(current_tokens) + len(sentence_tokens) > max_tokens:
+            if current_tokens:
+                segments.append(tokenizer.decode(current_tokens))
+            current_segment = words
+            current_tokens = sentence_tokens
+        else:
+            current_segment.extend(words)
+            current_tokens.extend(sentence_tokens)
+
+    if current_tokens:
+        segments.append(tokenizer.decode(current_tokens))
+
+    return segments
 
 
 def classify_segments(segments):
-    classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
-    classified_segments = []
-    
+    results = []
     for segment in segments:
         try:
-            if len(segment.split()) <= 512:  # Double-check to avoid errors
-                result = classifier(segment)
-                classified_segments.append(result)
-            else:
-                classified_segments.append({"error": f"Segment too long: {len(segment.split())} tokens"})
+            result = classifier(segment)
+            results.append(result)
         except Exception as e:
-            classified_segments.append({"error": str(e)})
+            results.append({"error": str(e), "segment": segment[:50]})  # Include a part of the segment to debug if needed
+    return results
 
-    return classified_segments
 
 
 @app.post("/process_document")
 async def process_document(request: TextRequest):
     try:
         processed_text = preprocess_text(request.text)
-        segments = segment_text(processed_text)
+        segments = robust_segment_text(processed_text)
         classified_segments = classify_segments(segments)
 
         return {