Update app.py

app.py

@@ -18,21 +18,26 @@ def classify_clauses(text_input):
     # Split input into clauses
     clauses = [clause.strip() for clause in text_input.strip().split('\n') if clause.strip()]
 
+    results = []
     scores = []
     for clause in clauses:
         result = classifier(clause, labels)
 
-        # Print the result to check the structure
-        print(f"Result for clause: '{clause}'")
-        print(result)  # Print the result to check its structure
-
-        # Ensure the 'scores' key exists in the result
+        # Extract scores and store the results
         if 'scores' in result:
             scores.append(result['scores'])
+            results.append({
+                'clause': clause,
+                'risk_score': result['scores'][0],  # High-risk score
+                'risk_level': labels[result['scores'].index(max(result['scores']))]  # Determine risk level
+            })
         else:
-            # Handle the case where 'scores' key is missing
-            print(f"Error: No scores found for clause: '{clause}'")
             scores.append([0, 0, 0])  # Default to no risk if 'scores' is missing
+            results.append({
+                'clause': clause,
+                'risk_score': 0,
+                'risk_level': 'low'
+            })
 
     scores_array = np.array(scores)
 
@@ -50,12 +55,12 @@ def classify_clauses(text_input):
     plt.ylabel("Clauses")
     plt.tight_layout()
 
-    # Save and return the plot
+    # Save and return the plot path
     plot_path = "heatmap.png"
     plt.savefig(plot_path)
     plt.close()
 
-    return plot_path
+    return results, plot_path
 
 # Function to analyze contract and generate results
 def analyze_contract(file):
@@ -65,7 +70,7 @@ def analyze_contract(file):
 
     # Split the text into clauses
     clauses = simple_clause_split(text)
-    results = classify_clauses("\n".join(clauses))  # Classify the clauses
+    results, heatmap_path = classify_clauses("\n".join(clauses))  # Classify the clauses
 
     # Calculate the overall risk score
     overall_score = sum(r['risk_score'] for r in results) / len(results) if results else 0
@@ -89,7 +94,7 @@ def analyze_contract(file):
         "overall_score": overall_score
     })
 
-    return f"Overall Risk Score: {overall_score:.2f}", highlight_output, report_path, results['heatmap']
+    return f"Overall Risk Score: {overall_score:.2f}", highlight_output, report_path, heatmap_path
 
 # Gradio interface
 demo = gr.Interface(