Update app.py

app.py

@@ -1,9 +1,8 @@
 import gradio as gr
-import matplotlib.pyplot as plt
-import seaborn as sns
 import numpy as np
-import os
+import plotly.graph_objects as go
 
+# Mock function to generate heatmap from contract analysis results
 def generate_heatmap(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
@@ -18,7 +17,7 @@ def generate_heatmap(results):
     risk_map = {'High': 3, 'Medium': 2, 'Low': 1}
     risk_values = [risk_map.get(risk, 0) for risk in risk_levels]
 
-    # Create a heatmap
+    # Create a heatmap using Plotly
     fig = go.Figure(data=go.Heatmap(
         z=[risk_values],  # The data for the heatmap
         x=clauses,  # Clause text as the x-axis labels
@@ -31,8 +30,20 @@ def generate_heatmap(results):
     fig.update_layout(
         title="Contract Risk Heatmap",
         xaxis_title="Clause",
-        yaxis_title="Risk Level"
+        yaxis_title="Risk Level",
+        xaxis=dict(tickmode='array', tickvals=range(len(clauses)), ticktext=[clause[:20] + '...' if len(clause) > 20 else clause for clause in clauses])  # Truncate clauses if too long
     )
 
     return fig.to_html(full_html=False)
 
+# Example usage with Gradio for testing the heatmap generator
+iface = gr.Interface(
+    fn=generate_heatmap,
+    inputs=gr.Textbox(label="Contract Text", placeholder="Enter the contract text here..."),
+    outputs=gr.HTML(),
+    title="Contract Risk Heatmap Generator",
+    description="Upload a contract and get a clause-level risk heatmap."
+)
+
+if __name__ == "__main__":
+    iface.launch()