no message

main.py

@@ -1,3 +1,4 @@
+import re
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.responses import JSONResponse
@@ -101,50 +102,27 @@ tokenizer = AutoTokenizer.from_pretrained("nsi319/legal-pegasus")
 model = AutoModelForSeq2SeqLM.from_pretrained("nsi319/legal-pegasus")
 
 
-class SummarizeRequest(BaseModel):
+class TextRequest(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)
+def preprocess_text(text: str) -> str:
+    # Normalize whitespace
+    text = re.sub(r'\s+', ' ', text.strip())
+
+    # Optional: Add additional preprocessing steps
+    # E.g., handling or stripping special characters, lowercasing, etc.
+    text = re.sub(r'[^\w\s]', '', text)  # Remove punctuation for simplicity
+
+    return text
 
 @app.post("/summarize")
-async def summarize_text(request: SummarizeRequest):
+async def summarize(request: TextRequest):
     try:
-        summarized_text = summarize_legal_text(request.text)
-        return JSONResponse(content={"summary": summarized_text})
+        processed_text = preprocess_text(request.text)
+        
+        return {"summary": processed_text}
+
     except Exception as e:
         print(f"Error during summarization: {e}")
         raise HTTPException(status_code=500, detail=str(e))