Update app.py

app.py

@@ -1,18 +1,17 @@
 import os
-import fitz  # PyMuPDF
-import docx
+import fitz  # PyMuPDF for PDF extraction
+import docx  # python-docx for DOCX extraction
 from sentence_transformers import SentenceTransformer, util
 import gradio as gr
-import torch
+import re
+from typing import List, Tuple, Dict
 import matplotlib.pyplot as plt
-import seaborn as sns
-from io import BytesIO
-import base64
+import numpy as np
+from collections import defaultdict
 
-# ----------------- Initialize model -----------------
+# Initialize the SentenceTransformer model
 model = SentenceTransformer('all-MiniLM-L6-v2')
 
-# ----------------- Text Extraction -----------------
 def extract_text_from_pdf(pdf_path):
     try:
         doc = fitz.open(pdf_path)
@@ -21,90 +20,184 @@ def extract_text_from_pdf(pdf_path):
             text += page.get_text()
         return text
     except Exception as e:
-        return f"Error extracting PDF: {str(e)}"
+        print(f"Error extracting text from PDF: {str(e)}")
+        return ""
 
 def extract_text_from_docx(docx_path):
     try:
         doc = docx.Document(docx_path)
-        return "\n".join([para.text for para in doc.paragraphs])
+        text = "\n".join([para.text for para in doc.paragraphs])
+        return text
     except Exception as e:
-        return f"Error extracting DOCX: {str(e)}"
-
-# ----------------- Chunk Similarity -----------------
-def chunk_text(text, chunk_size=5):
-    sentences = [s.strip() for s in text.split('.') if s.strip()]
-    chunks = ['. '.join(sentences[i:i+chunk_size]) for i in range(0, len(sentences), chunk_size)]
-    return chunks
+        print(f"Error extracting text from DOCX: {str(e)}")
+        return ""
 
-def get_similar_chunks(doc1, doc2, chunk_size=5, threshold=0.7):
-    chunks1 = chunk_text(doc1, chunk_size)
-    chunks2 = chunk_text(doc2, chunk_size)
-
-    embeddings1 = model.encode(chunks1, convert_to_tensor=True)
-    embeddings2 = model.encode(chunks2, convert_to_tensor=True)
+def preprocess_text(text: str) -> List[str]:
+    """Split text into sentences and clean them"""
+    # Split into sentences using regex
+    sentences = re.split(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?|\!)\s', text)
+    # Clean sentences
+    sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
+    return sentences
 
+def calculate_cosine_similarity(doc1: str, doc2: str) -> Tuple[float, List[Tuple[str, str, float]]]:
+    """Calculate similarity score and return similar sentence pairs"""
+    # Preprocess texts into sentences
+    sentences1 = preprocess_text(doc1)
+    sentences2 = preprocess_text(doc2)
+    
+    # Get embeddings for all sentences
+    embeddings1 = model.encode(sentences1, convert_to_tensor=True)
+    embeddings2 = model.encode(sentences2, convert_to_tensor=True)
+    
+    # Calculate cosine similarities between all sentence pairs
     cosine_similarities = util.pytorch_cos_sim(embeddings1, embeddings2)
-
+    
+    # Find the most similar sentences
     similar_pairs = []
-    for i, row in enumerate(cosine_similarities):
-        max_val, idx = row.max(0)
-        if max_val.item() >= threshold:
-            similar_pairs.append({
-                "doc1_chunk": chunks1[i],
-                "doc2_chunk": chunks2[idx],
-                "similarity": round(max_val.item(), 3)
-            })
-    return similar_pairs, cosine_similarities.cpu().numpy(), chunks1, chunks2
+    threshold = 0.7  # Similarity threshold for highlighting
+    
+    for i in range(len(sentences1)):
+        max_similarity = 0
+        best_match_idx = -1
+        
+        for j in range(len(sentences2)):
+            if cosine_similarities[i][j] > max_similarity:
+                max_similarity = cosine_similarities[i][j]
+                best_match_idx = j
+        
+        if max_similarity > threshold and best_match_idx != -1:
+            similar_pairs.append((sentences1[i], sentences2[best_match_idx], max_similarity.item()))
+    
+    # Calculate overall similarity
+    if len(sentences1) > 0 and len(sentences2) > 0:
+        # Use max similarity for each sentence and average
+        max_similarities1 = cosine_similarities.max(dim=1)[0]
+        max_similarities2 = cosine_similarities.max(dim=0)[0]
+        mean_similarity = (max_similarities1.mean() + max_similarities2.mean()) / 2.0
+        overall_similarity = mean_similarity.item()
+    else:
+        overall_similarity = 0.0
+    
+    return overall_similarity, similar_pairs
 
-# ----------------- Generate Heatmap -----------------
-def create_heatmap(sim_matrix, chunks1, chunks2):
-    plt.figure(figsize=(10,8))
-    sns.heatmap(sim_matrix, xticklabels=[f"C{i+1}" for i in range(len(chunks2))],
-                yticklabels=[f"C{i+1}" for i in range(len(chunks1))], cmap="viridis")
-    plt.xlabel("Document 2 Chunks")
-    plt.ylabel("Document 1 Chunks")
-    plt.title("Similarity Heatmap")
-    
-    buf = BytesIO()
-    plt.savefig(buf, format="png")
-    buf.seek(0)
-    encoded = base64.b64encode(buf.getvalue()).decode()
-    buf.close()
+def visualize_similarity(sentences1, sentences2, similarity_matrix):
+    """Create a heatmap visualization of sentence similarities"""
+    plt.figure(figsize=(10, 8))
+    plt.imshow(similarity_matrix, cmap='hot', interpolation='nearest')
+    plt.colorbar(label='Similarity Score')
+    plt.xlabel('Document 2 Sentences')
+    plt.ylabel('Document 1 Sentences')
+    plt.title('Sentence Similarity Heatmap')
+    plt.tight_layout()
+    plt.savefig('similarity_heatmap.png')
     plt.close()
-    return f"data:image/png;base64,{encoded}"
+    return 'similarity_heatmap.png'
+
+def group_similar_concepts(similar_pairs: List[Tuple[str, str, float]]) -> Dict[str, List[Tuple[str, str, float]]]:
+    """Group similar sentences by concept using keyword extraction"""
+    # Simple keyword-based grouping (could be enhanced with NLP techniques)
+    concept_groups = defaultdict(list)
+    
+    # Define some common concepts for SOPs
+    concepts = {
+        'research': ['research', 'study', 'investigate', 'experiment', 'methodology'],
+        'education': ['education', 'learn', 'course', 'degree', 'academic'],
+        'experience': ['experience', 'work', 'job', 'intern', 'position'],
+        'goals': ['goal', 'objective', 'aim', 'purpose', 'aspiration'],
+        'skills': ['skill', 'ability', 'proficient', 'expertise', 'capability']
+    }
+    
+    for sent1, sent2, score in similar_pairs:
+        matched_concept = 'other'
+        for concept, keywords in concepts.items():
+            if any(keyword in sent1.lower() for keyword in keywords) or \
+               any(keyword in sent2.lower() for keyword in keywords):
+                matched_concept = concept
+                break
+        concept_groups[matched_concept].append((sent1, sent2, score))
+    
+    return concept_groups
 
-# ----------------- Main Function -----------------
-def similarity_with_details(file1, file2, threshold=0.7):
+def similarity(file1, file2):
+    # Extract text based on file type
     text1 = extract_text_from_pdf(file1.name) if file1.name.endswith('.pdf') else extract_text_from_docx(file1.name)
     text2 = extract_text_from_pdf(file2.name) if file2.name.endswith('.pdf') else extract_text_from_docx(file2.name)
-
-    similar_pairs, sim_matrix, chunks1, chunks2 = get_similar_chunks(text1, text2, threshold=threshold)
     
-    if not similar_pairs:
-        return "No significant similarity found.", None
-
-    result = ""
-    for i, pair in enumerate(similar_pairs, 1):
-        result += f"### Similar Chunk {i} (Score: {pair['similarity']})\n"
-        result += f"**Doc1:** {pair['doc1_chunk']}\n"
-        result += f"**Doc2:** {pair['doc2_chunk']}\n\n"
-
-    heatmap_img = create_heatmap(sim_matrix, chunks1, chunks2)
-    return result, heatmap_img
+    # Calculate similarity and get similar pairs
+    overall_similarity, similar_pairs = calculate_cosine_similarity(text1, text2)
+    
+    # Group similar concepts
+    concept_groups = group_similar_concepts(similar_pairs)
+    
+    # Prepare detailed output
+    output_html = f"<h3>Overall Similarity Score: {overall_similarity:.2%}</h3>"
+    
+    if similar_pairs:
+        output_html += "<h4>Similar Content Found:</h4>"
+        
+        for concept, pairs in concept_groups.items():
+            if pairs:  # Only show concepts with matches
+                output_html += f"<h5>{concept.capitalize()}:</h5>"
+                for i, (sent1, sent2, score) in enumerate(pairs):
+                    output_html += f"""
+                    <div style="background-color: #f0f0f0; padding: 10px; margin: 5px; border-radius: 5px;">
+                        <p><b>Document 1:</b> {sent1}</p>
+                        <p><b>Document 2:</b> {sent2}</p>
+                        <p><b>Similarity:</b> {score:.2%}</p>
+                    </div>
+                    """
+    else:
+        output_html += "<p>No significant similarities found above the threshold.</p>"
+    
+    # Generate similarity heatmap if there are sentences
+    sentences1 = preprocess_text(text1)
+    sentences2 = preprocess_text(text2)
+    
+    if sentences1 and sentences2:
+        # Get embeddings for visualization
+        embeddings1 = model.encode(sentences1, convert_to_tensor=True)
+        embeddings2 = model.encode(sentences2, convert_to_tensor=True)
+        similarity_matrix = util.pytorch_cos_sim(embeddings1, embeddings2).cpu().numpy()
+        
+        # Generate and save heatmap
+        heatmap_path = visualize_similarity(sentences1, sentences2, similarity_matrix)
+        output_html += f'<h4>Similarity Heatmap:</h4><img src="/file={heatmap_path}" alt="Similarity Heatmap" style="max-width: 100%;">'
+    
+    return output_html
 
-# ----------------- Gradio Interface -----------------
-with gr.Blocks() as demo:
-    gr.Markdown("## 📄 Document Similarity Checker with Highlighted Chunks")
+# Create a Gradio interface with enhanced features
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # Document Similarity Checker with Detailed Analysis
+    Upload two documents to compare their content and identify specific similarities.
+    """)
+    
     with gr.Row():
-        file1 = gr.File(label="Upload Document 1")
-        file2 = gr.File(label="Upload Document 2")
-    threshold = gr.Slider(0, 1, value=0.7, step=0.05, label="Similarity Threshold")
-    output_text = gr.Markdown()
-    output_heatmap = gr.Image()
-    submit = gr.Button("Check Similarity")
-    
-    submit.click(fn=similarity_with_details, inputs=[file1, file2, threshold], outputs=[output_text, output_heatmap])
+        with gr.Column():
+            file1 = gr.File(label="Upload Document 1", file_types=[".pdf", ".docx"])
+            file2 = gr.File(label="Upload Document 2", file_types=[".pdf", ".docx"])
+            submit = gr.Button("Compare Documents", variant="primary")
+        
+        with gr.Column():
+            output = gr.HTML(label="Similarity Analysis Results")
+    
+    # Add examples for users to try
+    gr.Examples(
+        examples=[
+            [os.path.join(os.path.dirname(__file__), "sample1.pdf"), os.path.join(os.path.dirname(__file__), "sample2.pdf")],
+            [os.path.join(os.path.dirname(__file__), "sample1.docx"), os.path.join(os.path.dirname(__file__), "sample2.docx")]
+        ],
+        inputs=[file1, file2],
+        outputs=output,
+        fn=similarity,
+        cache_examples=False
+    )
+    
+    submit.click(fn=similarity, inputs=[file1, file2], outputs=output)
 
-# Run the Gradio app
+# Use the GRADIO_SERVER_PORT environment variable, default to 7860 if not set
 port = int(os.getenv('GRADIO_SERVER_PORT', 7860))
-demo.launch(server_port=port, server_name="0.0.0.0")
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=port)