Update app.py

app.py

@@ -1,50 +1,70 @@
-import torch
-from transformers import pipeline
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 import base64
 import pdfplumber
+from transformers import pipeline
+import torch
+from typing import List, Dict
+
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load BERT-based model for clause classification
+# Load the pre-trained BERT model for contract clause classification
 classifier = pipeline("text-classification", model="distilbert-base-uncased")
 
-# Function to extract text from PDF
-def extract_text_from_pdf(pdf_data):
+
+# Function to extract text from PDF file
+def extract_text_from_pdf(pdf_data: bytes) -> str:
     with pdfplumber.open(pdf_data) as pdf:
         text = ""
         for page in pdf.pages:
             text += page.extract_text() or ""
     return text
 
+
+# Define request body structure
+class ContractFile(BaseModel):
+    file: str  # Base64-encoded PDF file
+
+
 @app.post("/analyze_contract")
-async def analyze_contract(file: str):
-    # Decode base64 PDF
-    pdf_data = base64.b64decode(file)
-
-    # Extract text
-    contract_text = extract_text_from_pdf(pdf_data)
-
-    # Split into clauses (simple split for demo; use regex for production)
-    clauses = contract_text.split(". ")
-
-    # Analyze each clause
-    results = []
-    for clause in clauses:
-        if clause.strip():
-            result = classifier(clause)
-            risk_score = result[0]["score"] if result[0]["label"] == "POSITIVE" else 1 - result[0]["score"]
-            results.append({
-                "clause": clause,
-                "risk_level": "High" if risk_score > 0.7 else "Medium" if risk_score > 0.4 else "Low",
-                "risk_score": risk_score
-            })
-
-    # Calculate overall risk score
-    overall_score = sum(r["risk_score"] for r in results) / len(results) if results else 0
-
-    return {
-        "clauses": results,
-        "overall_score": overall_score
-    }
+async def analyze_contract(data: ContractFile):
+    try:
+        # Decode base64 PDF data
+        pdf_data = base64.b64decode(data.file)
+        
+        # Extract text from the PDF
+        contract_text = extract_text_from_pdf(pdf_data)
+        
+        # Split contract text into clauses (naive split by ".")
+        clauses = contract_text.split(". ")
+        
+        # Analyze each clause for risk level using the classifier
+        results = []
+        for clause in clauses:
+            if clause.strip():
+                result = classifier(clause)
+                risk_score = result[0]["score"] if result[0]["label"] == "POSITIVE" else 1 - result[0]["score"]
+                risk_level = "High" if risk_score > 0.7 else "Medium" if risk_score > 0.4 else "Low"
+                results.append({
+                    "clause": clause,
+                    "risk_level": risk_level,
+                    "risk_score": risk_score
+                })
+
+        # Calculate the overall risk score for the contract
+        overall_score = sum(r["risk_score"] for r in results) / len(results) if results else 0
+
+        return {
+            "clauses": results,
+            "overall_score": overall_score
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing the contract: {str(e)}")
+
+
+# Test root endpoint
+@app.get("/")
+async def read_root():
+    return {"message": "Welcome to the Contract Risk Heatmap Generator API!"}