Update app.py

app.py

@@ -1,4 +1,6 @@
 import os
+os.environ["TRANSFORMERS_NO_FAST"] = "1"  # Force use of slow tokenizers
+
 import io
 import torch
 import uvicorn
@@ -19,7 +21,7 @@ from pyngrok import ngrok
 from threading import Thread
 import time
 import uuid
-import subprocess  # Used for running ffmpeg commands
+import subprocess  # For running ffmpeg commands
 
 # ✅ Ensure compatibility with Google Colab
 try:
@@ -49,7 +51,7 @@ app.add_middleware(
 
 # ✅ Initialize document storage
 document_storage = {}
-chat_history = []  # ✅ Added global chat history
+chat_history = []  # Global chat history
 
 # ✅ Function to store document context by task ID
 def store_document_context(task_id, text):
@@ -68,26 +70,18 @@ def load_document_context(task_id):
 
 def fine_tune_cuad_model():
     """
-    Fine tunes a question-answering model on the CUAD (Contract Understanding Atticus Dataset)
-    for detailed clause extraction. This demo function uses one epoch for demonstration;
-    adjust training parameters as needed.
+    Fine tunes a QA model on the CUAD dataset for clause extraction.
+    This demo uses one epoch; adjust parameters as needed.
     """
     from datasets import load_dataset
     import numpy as np
-    # Optionally, load a metric (here we leave metrics out for brevity)
-    from transformers import Trainer, TrainingArguments
-    from transformers import AutoModelForQuestionAnswering
+    from transformers import Trainer, TrainingArguments, AutoModelForQuestionAnswering
 
     print("✅ Loading CUAD dataset for fine tuning...")
-    # Load the CUAD QA dataset (SQuAD-style) with custom code allowed
     dataset = load_dataset("theatticusproject/cuad-qa", trust_remote_code=True)
 
-    # Use the train split with a larger subset for production fine tuning
     if "train" in dataset:
-        # Select a larger subset for training, e.g., 1000 examples
         train_dataset = dataset["train"].select(range(1000))
-        
-        # For validation, you might select around 200 examples
         if "validation" in dataset:
             val_dataset = dataset["validation"].select(range(200))
         else:
@@ -99,12 +93,10 @@ def fine_tune_cuad_model():
 
     print("✅ Preparing training features...")
 
-    # Load a QA model and its tokenizer. Here we use deepset/roberta-base-squad2.
     tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
     model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
 
     def prepare_train_features(examples):
-        # Tokenize with question and context; use truncation only on the context.
         tokenized_examples = tokenizer(
             examples["question"],
             examples["context"],
@@ -153,11 +145,9 @@ def fine_tune_cuad_model():
     train_dataset = train_dataset.map(prepare_train_features, batched=True, remove_columns=train_dataset.column_names)
     val_dataset = val_dataset.map(prepare_train_features, batched=True, remove_columns=val_dataset.column_names)
 
-    # Set format for PyTorch QA training
     train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "start_positions", "end_positions"])
     val_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "start_positions", "end_positions"])
 
-    # For QA tasks, computing metrics can be more complex; here we skip metrics for brevity.
     training_args = TrainingArguments(
         output_dir="./fine_tuned_legal_qa",
         evaluation_strategy="steps",
@@ -170,7 +160,7 @@ def fine_tune_cuad_model():
         logging_steps=50,
         save_steps=100,
         load_best_model_at_end=True,
-        report_to=[]  # Disables wandb logging to avoid related issues
+        report_to=[]  # Disable wandb logging
     )
 
     print("✅ Starting fine tuning on CUAD QA dataset...")
@@ -203,8 +193,7 @@ try:
         nlp = spacy.load("en_core_web_sm")
     print("✅ Loading NLP models...")
 
-    # Updated summarizer initialization with a slow tokenizer
-    from transformers import AutoTokenizer
+    # Initialize summarizer with a slow tokenizer
     summarizer = pipeline(
         "summarization",
         model="nsi319/legal-pegasus",
@@ -213,14 +202,11 @@ try:
     )
     
     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)
-    speech_to_text = pipeline("automatic-speech-recognition",
-                             model="openai/whisper-medium",
-                             chunk_length_s=30,
-                             device_map="auto" if torch.cuda.is_available() else "cpu")
-
-    # ✅ Load or Fine Tune CUAD QA Model
+    ner_model = pipeline("ner", model="dslim/bert-base-NER", device=0 if torch.cuda.is_available() else -1)
+    speech_to_text = pipeline("automatic-speech-recognition", model="openai/whisper-medium", chunk_length_s=30,
+                              device_map="auto" if torch.cuda.is_available() else "cpu")
+
+    # Load or fine tune CUAD QA model
     if os.path.exists("fine_tuned_legal_qa"):
         print("✅ Loading fine-tuned CUAD QA model from fine_tuned_legal_qa...")
         cuad_tokenizer = AutoTokenizer.from_pretrained("fine_tuned_legal_qa")
@@ -239,7 +225,6 @@ except Exception as e:
     raise RuntimeError(f"Error loading models: {str(e)}")
 
 from transformers import pipeline
-
 qa_model = pipeline("question-answering", model="deepset/roberta-base-squad2")
 
 def legal_chatbot(user_input, context):
@@ -260,11 +245,10 @@ def extract_text_from_pdf(pdf_file):
         raise HTTPException(status_code=400, detail=f"PDF extraction failed: {str(e)}")
 
 def process_video_to_text(video_file_path):
-    """Extract audio from video using ffmpeg and convert to text."""
+    """Extracts audio from video using ffmpeg and converts to text."""
     try:
         print(f"Processing video file at {video_file_path}")
         temp_audio_path = os.path.join("temp", "extracted_audio.wav")
-        # Use ffmpeg command to extract audio from the video file
         cmd = [
             "ffmpeg", "-i", video_file_path, "-vn",
             "-acodec", "pcm_s16le", "-ar", "44100", "-ac", "2",
@@ -283,7 +267,7 @@ def process_video_to_text(video_file_path):
         raise HTTPException(status_code=400, detail=f"Video processing failed: {str(e)}")
 
 def process_audio_to_text(audio_file_path):
-    """Process audio file and convert to text."""
+    """Processes an audio file and converts it to text."""
     try:
         print(f"Processing audio file at {audio_file_path}")
         result = speech_to_text(audio_file_path)
@@ -429,7 +413,7 @@ def analyze_contract_clauses(text):
         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]  # Using start_logits for example
+        predictions = torch.sigmoid(outputs.start_logits).cpu().numpy()[0]
         for idx, confidence in enumerate(predictions):
             if confidence > 0.5 and idx < len(clause_types):
                 clauses_detected.append({"type": clause_types[idx], "confidence": float(confidence)})