Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 os.environ["TRANSFORMERS_NO_FAST"] = "1"  # Force use of slow tokenizers
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"   # Enable synchronous CUDA errors for debugging
 
 import io
 import torch
@@ -13,7 +14,7 @@ import numpy as np
 import json
 import tempfile
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form, BackgroundTasks
-from fastapi.responses import FileResponse, JSONResponse, HTMLResponse  # Added HTMLResponse
+from fastapi.responses import FileResponse, JSONResponse, HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
 from sentence_transformers import SentenceTransformer
@@ -34,12 +35,12 @@ from gensim import corpora, models
 # Global cache for analysis results based on file hash
 analysis_cache = {}
 
-# Ensure compatibility with Google Colab
+# Ensure compatibility with Google Colab (if applicable)
 try:
     from google.colab import drive
     drive.mount('/content/drive')
 except Exception:
-    pass  # Not running in Colab
+    pass
 
 # Ensure required directories exist
 os.makedirs("static", exist_ok=True)
@@ -64,16 +65,13 @@ app.add_middleware(
 document_storage = {}
 chat_history = []
 
-# Function to store document context by task ID
 def store_document_context(task_id, text):
     document_storage[task_id] = text
     return True
 
-# Function to load document context by task ID
 def load_document_context(task_id):
     return document_storage.get(task_id, "")
 
-# Utility to compute MD5 hash from file content
 def compute_md5(content: bytes) -> str:
     return hashlib.md5(content).hexdigest()
 
@@ -196,14 +194,13 @@ try:
         spacy.cli.download("en_core_web_sm")
         nlp = spacy.load("en_core_web_sm")
     print("✅ Loading NLP models...")
-    # Use Facebook's bart-large-cnn for summarization
+    from transformers import PegasusTokenizer  # Not used now since we're using bart-large-cnn
     summarizer = pipeline(
         "summarization",
         model="facebook/bart-large-cnn",
         tokenizer="facebook/bart-large-cnn",
         device=0 if torch.cuda.is_available() else -1
     )
-    # Optionally convert summarizer model to FP16 for faster inference on GPU
     if device == "cuda":
         summarizer.model.half()
 
@@ -344,7 +341,10 @@ def analyze_contract_clauses(text):
     # Create chunks of the text
     chunks = [text[i:i+max_length] for i in range(0, len(text), step) if i+step < len(text)]
     for chunk in chunks:
-        inputs = cuad_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512).to(device)
+        # Tokenize and move to GPU, then add safety clamp to avoid out-of-bound token indices
+        tokenized_inputs = cuad_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512)
+        inputs = {k: v.to(device) for k, v in tokenized_inputs.items()}
+        inputs["input_ids"] = torch.clamp(inputs["input_ids"], max=cuad_model.config.vocab_size - 1)
         with torch.no_grad():
             outputs = cuad_model(**inputs)
         predictions = torch.sigmoid(outputs.start_logits).cpu().numpy()[0]