Update app.py

app.py

@@ -1,63 +1,58 @@
+from transformers import pipeline
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
 import gradio as gr
-import pdfplumber
-import plotly.graph_objects as go
-
-# Function to extract text from PDF
-def extract_text_from_pdf(file_path):
-    with pdfplumber.open(file_path) as pdf:
-        text = ''
-        for page in pdf.pages:
-            text += page.extract_text()
-    return text
-
-# Function to generate heatmap from extracted text data (risk scores for clauses)
-def generate_heatmap(text):
-    # Assume you already have a mechanism to classify risk levels per clause
-    clauses = text.split('\n')  # Example: split by newlines to get clauses
-    risk_levels = ['High', 'Medium', 'Low']  # Placeholder risk levels for simplicity
-    
-    # Map risk levels to numeric values for plotting
-    risk_map = {'High': 3, 'Medium': 2, 'Low': 1}
-    risk_values = [risk_map.get(level, 1) for level in risk_levels * (len(clauses) // len(risk_levels))]
-    
-    # Create the heatmap
-    fig = go.Figure(data=go.Heatmap(
-        z=[risk_values],
-        x=clauses[:len(risk_values)],  # Only use the first N clauses
-        y=["Risk Levels"],
-        colorscale='YlOrRd',  # You can change the color scale
-        colorbar=dict(title="Risk Level")
-    ))
-
-    # Customize the layout
-    fig.update_layout(
-        title="Contract Risk Heatmap",
-        xaxis_title="Clause",
-        yaxis_title="Risk Level"
-    )
-
-    # Return the HTML representation of the heatmap
-    return fig.to_html(full_html=False)
-
-# Function to process the uploaded file and display the heatmap
-def process_contract(file):
-    # Extract text from PDF
-    text = extract_text_from_pdf(file.name)
-
-    # Generate heatmap
-    heatmap_html = generate_heatmap(text)
-    
-    return heatmap_html
-
-# Gradio Interface
+
+# Step 1: Load the classifier model
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+
+# Step 2: Function to process the input and generate the heatmap
+def generate_heatmap(file):
+    # Example clauses from a contract
+    clauses = [
+        "The contract may be terminated by either party with a 30-day notice.",
+        "The vendor is liable for all damages during service interruptions.",
+        "User data may be shared with third parties without notice.",
+        "Payment must be made within 15 days of invoice receipt.",
+        "No liability is accepted for indirect damages."
+    ]
+
+    # Step 3: Define candidate labels for risk
+    labels = ["high risk", "medium risk", "low risk"]
+
+    # Step 4: Generate scores for each clause
+    scores = []
+    for clause in clauses:
+        result = classifier(clause, labels)
+        scores.append(result['scores'])
+
+    # Convert to NumPy array for heatmap
+    scores_array = np.array(scores)
+
+    # Step 5: Plot heatmap
+    plt.figure(figsize=(10, 6))
+    sns.heatmap(scores_array, annot=True, xticklabels=labels, yticklabels=[f"Clause {i+1}" for i in range(len(clauses))], cmap="Reds")
+    plt.title("Contract Clause Risk Heatmap")
+    plt.xlabel("Risk Level")
+    plt.ylabel("Clauses")
+
+    # Save the heatmap as an image file
+    heatmap_path = "/mnt/data/contract_risk_heatmap.png"
+    plt.savefig(heatmap_path)
+    plt.close()
+
+    return heatmap_path
+
+# Step 6: Create Gradio interface
 iface = gr.Interface(
-    fn=process_contract,  # Function that processes the contract
-    inputs=gr.File(label="Upload Contract PDF"),  # Input: file upload
-    outputs=gr.HTML(label="Contract Risk Heatmap"),  # Output: HTML content (heatmap)
+    fn=generate_heatmap,
+    inputs=gr.File(label="Upload Contract Document"),
+    outputs=gr.Image(label="Risk Heatmap"),
     title="📜 Contract Risk Heatmap Generator",
-    description="Upload a contract and get clause-level risk scores with heatmap-style highlights."
+    description="Upload a contract and get a clause-level risk heatmap."
 )
 
-# Run the interface
+# Ensure the Gradio app starts properly
 if __name__ == "__main__":
     iface.launch()