Update app.py

app.py

@@ -16,7 +16,7 @@ from rouge_score import rouge_scorer
 from nltk.tokenize import sent_tokenize
 
 # Download NLTK data during the build process
-nltk.download('punkt_tab')
+nltk.download('punkt') # Changed 'punkt_tab' to 'punkt' for robustness
 nltk.download('stopwords')
 print(" NLTK data downloaded.")
 
@@ -70,11 +70,32 @@ class HybridLegalSummarizer:
             self.refinement_model = genai.GenerativeModel('models/gemini-2.5-flash')
 
     def get_legalbert_embedding(self, text):
+        # This function is now only used by the slow calculate_all_scores
         inputs = legalbert_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
         with torch.no_grad():
             outputs = legalbert_model(**inputs)
         return torch.mean(outputs.last_hidden_state, dim=1).squeeze().numpy()
 
+    # --- OPTIMIZATION 1: NEW BATCH FUNCTION ---
+    # This function processes a LIST of sentences at once
+    def get_legalbert_embeddings_batch(self, sentences_list):
+        if not sentences_list:
+            return np.array([])
+        # 1. Tokenize the whole batch
+        inputs = legalbert_tokenizer(
+            sentences_list, 
+            return_tensors="pt", 
+            padding=True, 
+            truncation=True, 
+            max_length=512
+        )
+        # 2. Run the model ONCE
+        with torch.no_grad():
+            outputs = legalbert_model(**inputs)
+        # 3. Get all embeddings
+        embeddings = torch.mean(outputs.last_hidden_state, dim=1)
+        return embeddings.cpu().numpy()
+
     def preprocess_text(self, text):
         text = re.sub(r'\s+', ' ', text)
         sentences = [s.strip() for s in sent_tokenize(text) if s.strip()]
@@ -89,12 +110,18 @@ class HybridLegalSummarizer:
     def generate_extractive_draft(self, text, max_words=200):
         sentences = sent_tokenize(text)
         if not sentences: return ""
-        # USE THE BATCH FUNCTION
+        
+        # --- OPTIMIZATION 2: BATCHED & VECTORIZED ---
+        # 1. Get all embeddings at once (replaces a for loop)
         sentence_embeddings = self.get_legalbert_embeddings_batch(sentences)
         if sentence_embeddings.size == 0: return ""
         
         centroid = np.mean(sentence_embeddings, axis=0)
-        scores = [cosine_similarity(emb.reshape(1, -1), centroid.reshape(1, -1))[0][0] for emb in sentence_embeddings]
+        
+        # 2. Get all scores at once (replaces another for loop)
+        scores = cosine_similarity(sentence_embeddings, centroid.reshape(1, -1)).flatten()
+        # --- END OPTIMIZATION ---
+
         ranked_indices = np.argsort(scores)[::-1]
         selected = []
         current_count = 0
@@ -111,12 +138,17 @@ class HybridLegalSummarizer:
     def generate_rag_draft(self, text, user_query, max_words=600):
         sentences = sent_tokenize(text)
         if not sentences: return ""
-        # USE THE BATCH FUNCTION FOR BOTH
+
+        # --- OPTIMIZATION 2: BATCHED & VECTORIZED ---
+        # 1. Get query and sentence embeddings at once
         query_embedding = self.get_legalbert_embeddings_batch([user_query])[0]
         sentence_embeddings = self.get_legalbert_embeddings_batch(sentences)
         if sentence_embeddings.size == 0: return ""
-            
-        scores = [cosine_similarity(emb.reshape(1, -1), query_embedding.reshape(1, -1))[0][0] for emb in sentence_embeddings]
+
+        # 2. Get all scores at once (replaces a for loop)
+        scores = cosine_similarity(sentence_embeddings, query_embedding.reshape(1, -1)).flatten()
+        # --- END OPTIMIZATION ---
+
         ranked_indices = np.argsort(scores)[::-1]
         selected = []
         current_count = 0
@@ -141,35 +173,16 @@ class HybridLegalSummarizer:
             return response.text.strip()
         except Exception as e:
             return f"Refinement failed. Error: {e}"
-    # NEW BATCH FUNCTION TO ADD
-    def get_legalbert_embeddings_batch(self, sentences_list):
-        if not sentences_list:
-            return np.array([])
-            
-        # 1. Tokenize the whole batch at once
-        inputs = legalbert_tokenizer(
-            sentences_list, 
-            return_tensors="pt", 
-            padding=True, 
-            truncation=True, 
-            max_length=512
-        )
-        
-        # 2. Run the model ONCE on the whole batch
-        with torch.no_grad():
-            outputs = legalbert_model(**inputs)
-        
-        # 3. Use "mean pooling" to get the sentence embeddings
-        # This averages the tokens for each sentence to get a single vector
-        embeddings = torch.mean(outputs.last_hidden_state, dim=1)
-        return embeddings.cpu().numpy()
-        
+
     def calculate_all_scores(self, original_text, summary):
         if not summary or not original_text: return {}, 0.0, 0.0
         rouge = self.rouge_scorer.score(original_text, summary)
         rouge_scores = {"rouge1": rouge['rouge1'].fmeasure, "rouge2": rouge['rouge2'].fmeasure, "rougeL": rouge['rougeL'].fmeasure}
         
-        # --- TEMPORARILY DISABLED FOR SPEED ---
+        # --- OPTIMIZATION 3: DISABLED SLOW METRICS ---
+        # The following lines are too slow for a live demo as they run
+        # the model on the *entire* text.
+        
         # orig_emb = self.get_legalbert_embedding(original_text).reshape(1, -1)
         # sum_emb = self.get_legalbert_embedding(summary).reshape(1, -1)
         # consistency = cosine_similarity(orig_emb, sum_emb)[0][0]
@@ -177,10 +190,10 @@ class HybridLegalSummarizer:
         # sum_kw = self.extract_legal_terms(summary)
         # coverage = (len(orig_kw.intersection(sum_kw)) / len(orig_kw) * 100) if orig_kw else 0
         
-        # Just return dummy values for the disabled metrics
+        # Return dummy values for a fast demo.
         consistency = 0.0
         coverage = 0.0
-        # --- END TEMPORARY FIX ---
+        # --- END OPTIMIZATION ---
 
         return rouge_scores, consistency, coverage
 
@@ -214,11 +227,14 @@ def process_document(pdf_file, mode, word_limit, query):
     else:
         return "Error: Invalid mode selected.", ""
 
+    # This part is now fast because calculate_all_scores is fast
     final_rouge, final_consistency, final_coverage = summarizer.calculate_all_scores(cleaned_text, final_output)
+    
+    # Updated metrics string to show which values are disabled
     metrics_str = (
         f"ROUGE Scores: R1: {final_rouge.get('rouge1', 0):.3f}, R2: {final_rouge.get('rouge2', 0):.3f}, RL: {final_rouge.get('rougeL', 0):.3f}\n"
-        f" Factual Consistency (Semantic Similarity): {final_consistency:.3f}\n"
-        f" Legal Keyword Coverage: {final_coverage:.1f}%\n"
+        f" Factual Consistency (Semantic Similarity): {final_consistency:.3f} (Disabled for demo speed)\n"
+        f" Legal Keyword Coverage: {final_coverage:.1f}% (Disabled for demo speed)\n"
         f"Words in Output: {len(final_output.split())}"
     )
     return final_output, metrics_str