Update app.py

app.py

@@ -5,38 +5,34 @@ import numpy as np
 import os
 
 def generate_heatmap(results):
-    # Check the structure of the results
-    print("Results:", results)
+    # If results is a string, process it into a list of dicts (or just use the text as is)
+    if isinstance(results, str):  # Check if results is a string
+        # For testing, we will mock the results with dummy values if it's just text
+        results = [{"clause": results, "risk_level": "High"}]
 
-    # If the results are strings (e.g., just the clauses)
-    if isinstance(results, list) and isinstance(results[0], str):
-        clauses = results  # Directly use clauses
-        # For simplicity, assume all clauses are "high risk" here for testing purposes
-        risk_values = [3 for _ in clauses]  # Replace with actual risk assessment logic
-    else:
-        # Assuming results are in the format [{'clause': ..., 'risk_level': ...}, ...]
-        clauses = [r['clause'] for r in results]  # Extract clause text
-        risk_levels = {"High": 3, "Medium": 2, "Low": 1}
-        risk_values = [risk_levels.get(r['risk_level'], 1) for r in results]  # Map risk level to value
+    # Assuming results is now a list of dictionaries with keys 'clause' and 'risk_level'
+    clauses = [r['clause'] for r in results]  # Extract clause text
+    risk_levels = [r['risk_level'] for r in results]  # Extract risk levels
+    
+    # Map risk levels to numeric values for plotting
+    risk_map = {'High': 3, 'Medium': 2, 'Low': 1}
+    risk_values = [risk_map.get(risk, 0) for risk in risk_levels]
 
-    # Plot heatmap
-    fig = plt.figure(figsize=(10, 6))
-    sns.heatmap([risk_values], annot=True, xticklabels=clauses, yticklabels=["Risk Levels"], cmap="YlOrRd")
+    # Create a heatmap
+    fig = go.Figure(data=go.Heatmap(
+        z=[risk_values],  # The data for the heatmap
+        x=clauses,  # Clause text as the x-axis labels
+        y=["Risk Levels"],  # Y-axis representing risk levels
+        colorscale='YlOrRd',  # You can change the color scale
+        colorbar=dict(title="Risk Level")
+    ))
 
-    # Save heatmap image
-    heatmap_path = os.path.join(os.getcwd(), 'contract_risk_heatmap.png')
-    plt.savefig(heatmap_path)
+    # Add titles and labels to the heatmap
+    fig.update_layout(
+        title="Contract Risk Heatmap",
+        xaxis_title="Clause",
+        yaxis_title="Risk Level"
+    )
 
-    return heatmap_path
+    return fig.to_html(full_html=False)
 
-# Gradio interface
-iface = gr.Interface(
-    fn=generate_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."
-)
-
-# Start the app
-iface.launch()