Update app.py

app.py

@@ -1,6 +1,5 @@
 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
@@ -35,12 +34,12 @@ from gensim import corpora, models
 # Global cache for analysis results based on file hash
 analysis_cache = {}
 
-# Ensure compatibility with Google Colab (if applicable)
+# Ensure compatibility with Google Colab
 try:
     from google.colab import drive
     drive.mount('/content/drive')
 except Exception:
-    pass
+    pass  # Not running in Colab
 
 # Ensure required directories exist
 os.makedirs("static", exist_ok=True)
@@ -65,13 +64,16 @@ 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()
 
@@ -194,15 +196,16 @@ try:
         spacy.cli.download("en_core_web_sm")
         nlp = spacy.load("en_core_web_sm")
     print("✅ Loading NLP models...")
-    from transformers import PegasusTokenizer  # Not used now since we're using bart-large-cnn
+    # Use Facebook's bart-large-cnn for summarization
     summarizer = pipeline(
         "summarization",
         model="facebook/bart-large-cnn",
         tokenizer="facebook/bart-large-cnn",
         device=0 if torch.cuda.is_available() else -1
     )
-    if device == "cuda":
-        summarizer.model.half()
+    # Removed FP16 conversion for summarizer to avoid CUDA errors
+    # if device == "cuda":
+    #     summarizer.model.half()
 
     embedding_model = SentenceTransformer("all-mpnet-base-v2", device=device)
     ner_model = pipeline("ner", model="dslim/bert-base-NER", device=0 if torch.cuda.is_available() else -1)
@@ -341,10 +344,7 @@ 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:
-        # 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)
+        inputs = cuad_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512).to(device)
         with torch.no_grad():
             outputs = cuad_model(**inputs)
         predictions = torch.sigmoid(outputs.start_logits).cpu().numpy()[0]