Update app.py

app.py

@@ -1,65 +1,69 @@
-import os
+import streamlit as st
+import requests
 import json
-from transformers import BertTokenizer, BertForSequenceClassification
-import torch
-from flask import Flask, request, jsonify
-from pathlib import Path
-from datetime import datetime
+from transformers import pipeline
 
-# Initialize Flask app
-app = Flask(__name__)
+# Initialize the BERT-based NLP pipeline
+model_name = "your-huggingface-model-name"  # Replace this with your model
+nlp_pipeline = pipeline("ner", model=model_name)
 
-# Load pre-trained model and tokenizer
-MODEL_PATH = "path/to/your/model"  # Update with your Hugging Face model path
-tokenizer = BertTokenizer.from_pretrained(MODEL_PATH)
-model = BertForSequenceClassification.from_pretrained(MODEL_PATH)
-
-# Function to process contract text and classify clauses
-def classify_clause(contract_text):
-    inputs = tokenizer(contract_text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class = torch.argmax(logits, dim=-1).item()  # Get predicted class (risk tag)
-    
-    # Define risk labels (assuming 3 risk levels: low, medium, high)
-    risk_labels = ["low", "medium", "high"]
-    predicted_risk = risk_labels[predicted_class]
-
-    # Get confidence score (softmax output)
-    softmax = torch.nn.Softmax(dim=-1)
-    confidence = softmax(logits).squeeze().tolist()[predicted_class]
-
-    return {"predicted_risk": predicted_risk, "confidence_score": confidence}
-
-# Define route to handle file uploads
-@app.route("/upload_contract", methods=["POST"])
-def upload_contract():
-    # Extract file from the request
-    if 'file' not in request.files:
-        return jsonify({"error": "No file part"}), 400
+# Function to analyze contract text
+def analyze_contract(contract_text):
+    # Run the contract through the NLP pipeline
+    results = nlp_pipeline(contract_text)
     
-    file = request.files['file']
+    # Parse and score clauses (this is a simplified version)
+    risk_score = 0
+    high_risk_clauses = []
     
-    if file.filename == '':
-        return jsonify({"error": "No selected file"}), 400
+    for result in results:
+        # This assumes 'labels' are risk-related; adjust as per model output
+        if result['label'] in ["PENALTY", "OBLIGATION", "DELAY"]:  # Customize as per your model's tags
+            high_risk_clauses.append(result['word'])
+            risk_score += 10  # Example scoring logic, modify as needed
     
-    contract_text = file.read().decode('utf-8')  # Assuming the file is a text-based contract
-    
-    # Classify the contract text
-    result = classify_clause(contract_text)
-
-    # Prepare JSON response
-    response_data = {
-        "contract_title": "Sample Contract",  # Placeholder, can be parsed from the file
-        "overall_risk_score": result["predicted_risk"],  # Risk classification
-        "high_risk_clauses": ["Termination Clause", "Penalty Clause"],  # Example (this should be dynamically extracted)
-        "risk_map_url": "https://example.com/risk_map",  # Placeholder (use actual URL for visualization)
-        "evaluation_date": datetime.now().strftime("%Y-%m-%d")
+    return {
+        "high_risk_clauses": high_risk_clauses,
+        "risk_score": risk_score
     }
 
-    # Return response as JSON
-    return jsonify(response_data)
+# Streamlit UI
+st.title("Contract Risk Analyzer")
+
+# File upload
+contract_file = st.file_uploader("Upload Contract", type=["pdf", "docx", "txt"])
+
+if contract_file is not None:
+    contract_text = ""
+    if contract_file.type == "application/pdf":
+        import PyPDF2
+        # Read PDF
+        pdf_reader = PyPDF2.PdfReader(contract_file)
+        for page in pdf_reader.pages:
+            contract_text += page.extract_text()
+    elif contract_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+        import docx
+        # Read DOCX
+        doc = docx.Document(contract_file)
+        for para in doc.paragraphs:
+            contract_text += para.text
+    elif contract_file.type == "text/plain":
+        contract_text = contract_file.read().decode("utf-8")
+
+    # Analyze the contract text
+    if contract_text:
+        analysis_results = analyze_contract(contract_text)
+        
+        # Display the high-risk clauses and risk score
+        st.subheader("High Risk Clauses")
+        st.write(", ".join(analysis_results["high_risk_clauses"]))
+        
+        st.subheader("Overall Risk Score")
+        st.write(analysis_results["risk_score"])
+        
+        # Generate the risk heatmap (simplified here, you might want a more complex rendering)
+        st.subheader("Risk Heatmap")
+        st.write(f"Risk Score: {analysis_results['risk_score']}")
+        # Visualize as per your design (here we can display a simple score)
 
-if __name__ == "__main__":
-    app.run(debug=True, host="0.0.0.0", port=5000)
+        # Here you could add logic to save the results to Salesforce or other systems