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') # Changed 'punkt_tab' to 'punkt' for robustness
+nltk.download('punkt_tab')
 nltk.download('stopwords')
 print(" NLTK data downloaded.")
 
@@ -70,32 +70,11 @@ 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()]
@@ -110,18 +89,9 @@ class HybridLegalSummarizer:
     def generate_extractive_draft(self, text, max_words=200):
         sentences = sent_tokenize(text)
         if not sentences: return ""
-        
-        # --- 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 ""
-        
+        sentence_embeddings = np.array([self.get_legalbert_embedding(sent) for sent in sentences])
         centroid = np.mean(sentence_embeddings, axis=0)
-        
-        # 2. Get all scores at once (replaces another for loop)
-        scores = cosine_similarity(sentence_embeddings, centroid.reshape(1, -1)).flatten()
-        # --- END OPTIMIZATION ---
-
+        scores = [cosine_similarity(emb.reshape(1, -1), centroid.reshape(1, -1))[0][0] for emb in sentence_embeddings]
         ranked_indices = np.argsort(scores)[::-1]
         selected = []
         current_count = 0
@@ -138,17 +108,9 @@ class HybridLegalSummarizer:
     def generate_rag_draft(self, text, user_query, max_words=600):
         sentences = sent_tokenize(text)
         if not sentences: return ""
-
-        # --- 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 ""
-
-        # 2. Get all scores at once (replaces a for loop)
-        scores = cosine_similarity(sentence_embeddings, query_embedding.reshape(1, -1)).flatten()
-        # --- END OPTIMIZATION ---
-
+        query_embedding = self.get_legalbert_embedding(user_query)
+        sentence_embeddings = np.array([self.get_legalbert_embedding(sent) for sent in sentences])
+        scores = [cosine_similarity(emb.reshape(1, -1), query_embedding.reshape(1, -1))[0][0] for emb in sentence_embeddings]
         ranked_indices = np.argsort(scores)[::-1]
         selected = []
         current_count = 0
@@ -178,23 +140,12 @@ class HybridLegalSummarizer:
         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}
-        
-        # --- 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]
-        # orig_kw = self.extract_legal_terms(original_text)
-        # sum_kw = self.extract_legal_terms(summary)
-        # coverage = (len(orig_kw.intersection(sum_kw)) / len(orig_kw) * 100) if orig_kw else 0
-        
-        # Return dummy values for a fast demo.
-        consistency = 0.0
-        coverage = 0.0
-        # --- END OPTIMIZATION ---
-
+        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]
+        orig_kw = self.extract_legal_terms(original_text)
+        sum_kw = self.extract_legal_terms(summary)
+        coverage = (len(orig_kw.intersection(sum_kw)) / len(orig_kw) * 100) if orig_kw else 0
         return rouge_scores, consistency, coverage
 
 summarizer = HybridLegalSummarizer()
@@ -227,14 +178,11 @@ 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} (Disabled for demo speed)\n"
-        f" Legal Keyword Coverage: {final_coverage:.1f}% (Disabled for demo speed)\n"
+        f" Factual Consistency (Semantic Similarity): {final_consistency:.3f}\n"
+        f" Legal Keyword Coverage: {final_coverage:.1f}%\n"
         f"Words in Output: {len(final_output.split())}"
     )
     return final_output, metrics_str