no message

main.py

@@ -100,33 +100,49 @@ def split_text_by_tokens(text, max_tokens=1024):
 tokenizer = AutoTokenizer.from_pretrained("nsi319/legal-pegasus")
 model = AutoModelForSeq2SeqLM.from_pretrained("nsi319/legal-pegasus")
 
-def summarize_legal_text(text):
-    # Ensure the text is within the maximum length limit for the model
-    inputs = tokenizer.encode(text, return_tensors='pt', max_length=1024, truncation=True)
-    
-    # Generate summary
-    summary_ids = model.generate(
-        inputs,
-        num_beams=5,
-        no_repeat_ngram_size=3,
-        length_penalty=1.0,
-        min_length=150,
-        max_length=1000,
-        early_stopping=True
-    )
-    
-    # Decode generated tokens to a string
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return summary
-
 
 class SummarizeRequest(BaseModel):
     text: str
 
+def chunk_text(text, max_length=1024):
+    """Split the text into manageable parts for the model to handle."""
+    words = text.split()
+    current_chunk = ""
+    chunks = []
+
+    for word in words:
+        if len(tokenizer.encode(current_chunk + word)) < max_length:
+            current_chunk += word + ' '
+        else:
+            chunks.append(current_chunk.strip())
+            current_chunk = word + ' '
+    chunks.append(current_chunk.strip())  # Add the last chunk
+    return chunks
+
+def summarize_legal_text(text):
+    """Generate summaries for each chunk and combine them."""
+    chunks = chunk_text(text, max_length=900)  # A bit less than 1024 to be safe
+    all_summaries = []
+
+    for chunk in chunks:
+        inputs = tokenizer.encode(chunk, return_tensors='pt', max_length=1024, truncation=True)
+        summary_ids = model.generate(
+            inputs,
+            num_beams=5,
+            no_repeat_ngram_size=3,
+            length_penalty=1.0,
+            min_length=150,
+            max_length=300,  # You can adjust this based on your needs
+            early_stopping=True
+        )
+        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+        all_summaries.append(summary)
+
+    return " ".join(all_summaries)
+
 @app.post("/summarize")
 async def summarize_text(request: SummarizeRequest):
     try:
-        # Use the newly defined summarization function
         summarized_text = summarize_legal_text(request.text)
         return JSONResponse(content={"summary": summarized_text})
     except Exception as e: