no message

main.py

@@ -108,63 +108,70 @@ def reduce_tokens(text: str):
     token_count = len(reduced_doc)
     return reduced_text, token_count
 
-def segment_text(text: str, max_tokens=500):  # Slightly less than 512 for safety
-    # Use spaCy to divide the document into sentences
+def segment_text(text: str, max_tokens=500):  # Setting a conservative limit below 512
     doc = nlp(text)
-    sentences = [sent.text.strip() for sent in doc.sents]
-
     segments = []
     current_segment = []
     current_length = 0
-    
-    for sentence in sentences:
-        sentence_words = sentence.split()
-        sentence_length = len(sentence_words)
-        
-        # If sentence exceeds max_tokens, split it further
+
+    for sent in doc.sents:
+        sentence = sent.text.strip()
+        sentence_length = len(sentence.split())  # Counting words for simplicity
+
         if sentence_length > max_tokens:
-            parts = split_into_parts(sentence, max_tokens)
-            segments.extend(parts)  # Add split parts directly to segments
-            continue
-        
-        if current_length + sentence_length > max_tokens:
+            # Split long sentences into smaller chunks if a single sentence exceeds max_tokens
+            words = sentence.split()
+            while words:
+                part = ' '.join(words[:max_tokens])
+                segments.append(part)
+                words = words[max_tokens:]
+        elif current_length + sentence_length > max_tokens:
             segments.append(' '.join(current_segment))
             current_segment = [sentence]
             current_length = sentence_length
         else:
             current_segment.append(sentence)
             current_length += sentence_length
-    
-    if current_segment:  # Add the last segment if any
-        segments.append(' '.join(current_segment))
-    
-    return segments
 
-def split_into_parts(text, max_tokens):
-    words = text.split()
-    parts = []
-    for i in range(0, len(words), max_tokens):
-        part = " ".join(words[i:i + max_tokens])
-        parts.append(part)
-    return parts
+    if current_segment:  # Add the last segment
+        segments.append(' '.join(current_segment))
 
+    return segments
 
 
 classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
 
 
 def classify_segments(segments):
-    results = []
+    classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
+    classified_segments = []
+    
     for segment in segments:
         try:
-            if len(segment.split()) <= 512:  # Ensure segment is within the limit
+            if len(segment.split()) <= 512:  # Double-check to avoid errors
                 result = classifier(segment)
-                results.append(result)
+                classified_segments.append(result)
             else:
-                results.append({"error": f"Segment too long: {len(segment.split())} tokens"})
+                classified_segments.append({"error": f"Segment too long: {len(segment.split())} tokens"})
         except Exception as e:
-            results.append({"error": str(e)})
-    return results
+            classified_segments.append({"error": str(e)})
+
+    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)
+        classified_segments = classify_segments(segments)
+
+        return {
+            "classified_segments": classified_segments
+        }
+    except Exception as e:
+        print(f"Error during document processing: {e}")
+        raise HTTPException(status_code=500, detail=str(e))