no message

main.py

@@ -14,6 +14,7 @@ import sentencepiece
 from transformers import pipeline, AutoTokenizer,AutoModelForSequenceClassification,AutoModel
 import spacy
 import numpy as np
+import torch
 
 
 nltk.data.path.append(os.getenv('NLTK_DATA'))
@@ -80,25 +81,39 @@ async def generate_text(item: Item):
     # Stream response back to the client
     return StreamingResponse(generate_stream(item), media_type="application/x-ndjson")
 
-
-
-
 # Define request model
 class TextRequest(BaseModel):
-    text: List[str]  # Expect a list of text segments
+    text: str  # Single string of long text
 
 # Load Longformer model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")
-model = AutoModelForSequenceClassification.from_pretrained("allenai/longformer-base-4096")
+model = AutoModel.from_pretrained("allenai/longformer-base-4096")
 
-# Endpoint to process the document and return embeddings for each segment
+# Endpoint to process the document and return embeddings
 @app.post("/process_document")
 async def process_document(request: TextRequest):
     try:
+        # Split the text into segments that fit within the model's max input size
+        max_length = 4096  # Maximum token length for Longformer
+        words = request.text.split()
+        tokens = tokenizer.encode(request.text, add_special_tokens=True)
+        input_ids = []
+        current_chunk = []
+
+        for token in tokens:
+            if len(current_chunk) + len(tokenizer.convert_ids_to_tokens([token])) < max_length:
+                current_chunk.append(token)
+            else:
+                input_ids.append(current_chunk)
+                current_chunk = [token]
+
+        if current_chunk:
+            input_ids.append(current_chunk)  # Add the last chunk if any
+
+        # Generate embeddings for each segment
         embeddings_list = []
-        for text_segment in request.text:
-            # Process each segment individually
-            inputs = tokenizer(text_segment, return_tensors="pt", padding=True, truncation=True, max_length=4096)
+        for ids in input_ids:
+            inputs = {'input_ids': torch.tensor(ids).unsqueeze(0)}  # Batch size 1
             outputs = model(**inputs)
             embeddings = outputs.last_hidden_state.mean(dim=1).detach().numpy()
             embeddings_list.append(embeddings.tolist())  # Store embeddings for each segment