no message

main.py

@@ -8,7 +8,7 @@ from typing import Generator
 import json  # Asegúrate de que esta línea esté al principio del archivo
 import nltk
 import os
-from transformers import pipeline, BartTokenizer
+from transformers import pipeline, AutoModelForSeq2SeqLM
 
 
 nltk.data.path.append(os.getenv('NLTK_DATA'))
@@ -18,9 +18,7 @@ app = FastAPI()
 # Initialize the InferenceClient with your model
 client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
 
-# Assuming you've initialized the tokenizer and model for BART
-tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
-summarizer = pipeline("summarization", model="facebook/bart-large")
+
 
 # summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
 
@@ -96,31 +94,43 @@ def split_text_by_tokens(text, max_tokens=1024):
     print("Tokenization complete.")
     return chunks, token_counts
 
-def summarize_large_text(text):
-    chunks, token_counts = split_text_by_tokens(text, max_tokens=1024 - 10)  # Slight buffer to avoid edge cases
-    summaries = []
-    print("Starting summarization of chunks...")
+# Load the tokenizer and model from Hugging Face Hub
+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=9,
+        no_repeat_ngram_size=3,
+        length_penalty=2.0,
+        min_length=150,
+        max_length=250,
+        early_stopping=True
+    )
+    
+    # Decode generated tokens to a string
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
 
-    for chunk in chunks:
-        print(f"Summarizing chunk: {chunk[:50]}...")  # Print the first 50 characters of the chunk
-        # Perform summarization on the chunk
-        summary = summarizer(chunk, max_length=500, min_length=100, do_sample=False)
-        if summary:
-            summaries.append(summary[0]['summary_text'])
-            print(f"Summary: {summary[0]['summary_text'][:50]}")  # Print the first 50 characters of the summary
-    combined_summary = ' '.join(summaries)
-    return combined_summary
 
+class SummarizeRequest(BaseModel):
+    text: str
 
 @app.post("/summarize")
 async def summarize_text(request: SummarizeRequest):
     try:
-        summarized_text = summarize_large_text(request.text)
+        # Use the newly defined summarization function
+        summarized_text = summarize_legal_text(request.text)
         return JSONResponse(content={"summary": summarized_text})
     except Exception as e:
-        print(f"Error during tokenization: {e}")
+        print(f"Error during summarization: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
-    
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)