Update app.py

app.py

@@ -1,58 +1,33 @@
-from transformers import pipeline
+import gradio as gr
 import matplotlib.pyplot as plt
 import seaborn as sns
 import numpy as np
-import gradio as gr
-
-# 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'])
+import os
 
-    # Convert to NumPy array for heatmap
-    scores_array = np.array(scores)
+def generate_heatmap(results):
+    # Sample risk levels
+    risk_levels = {"High": 3, "Medium": 2, "Low": 1}
+    
+    clauses = [r['clause'] for r in results]
+    risk_values = [risk_levels.get(r['risk_level'], 1) for r in results]
 
-    # 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")
+    fig = plt.figure(figsize=(10, 6))
+    sns.heatmap([risk_values], annot=True, xticklabels=clauses, yticklabels=["Risk Levels"], cmap="YlOrRd")
 
-    # Save the heatmap as an image file
-    heatmap_path = "/mnt/data/contract_risk_heatmap.png"
+    # Save heatmap image to current directory
+    heatmap_path = os.path.join(os.getcwd(), 'contract_risk_heatmap.png')
     plt.savefig(heatmap_path)
-    plt.close()
 
     return heatmap_path
 
-# Step 6: Create Gradio interface
+# Gradio interface
 iface = gr.Interface(
     fn=generate_heatmap,
-    inputs=gr.File(label="Upload Contract Document"),
-    outputs=gr.Image(label="Risk Heatmap"),
+    inputs=gr.File(label="Upload Contract PDF"),
+    outputs=gr.File(label="Download Heatmap Image"),
     title="📜 Contract Risk Heatmap Generator",
     description="Upload a contract and get a clause-level risk heatmap."
 )
 
-# Ensure the Gradio app starts properly
-if __name__ == "__main__":
-    iface.launch()
+# Start the app
+iface.launch()