Update app.py

app.py

@@ -1,166 +1,133 @@
-import gradio as gr
-import PyPDF2
-import nltk
+import streamlit as st
+import pandas as pd
+import numpy as np
+from transformers import BertTokenizer, BertForSequenceClassification
+from simple_salesforce import Salesforce
+import torch
+from PyPDF2 import PdfReader
+import re
 import seaborn as sns
 import matplotlib.pyplot as plt
-from reportlab.lib.pagesizes import letter
-from reportlab.pdfgen import canvas
-import json
-import os
-from io import BytesIO
-import numpy as np
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Download NLTK data
-nltk.download('punkt')
-
-# Clause types and risk scoring logic
-CLAUSE_TYPES = ["penalty", "obligation", "delay"]
-RISK_WEIGHTS = {"penalty": 0.8, "obligation": 0.5, "delay": 0.6}
-
-# Keyword-based heuristic for clause classification
-KEYWORD_MAP = {
-    "penalty": ["penalty", "fee", "fine", "charge", "incur"],
-    "obligation": ["shall", "must", "obligated", "required", "responsible"],
-    "delay": ["delay", "late", "beyond", "postpone", "deferred"]
-}
 
+# Salesforce connection
+def connect_to_salesforce():
+    sf = Salesforce(
+        username='your_username',
+        password='your_password',
+        security_token='your_security_token',
+        domain='login'  # or 'test' for sandbox
+    )
+    return sf
+
+# Extract text from PDF
 def extract_text_from_pdf(pdf_file):
-    """Extract text from uploaded PDF file."""
-    try:
-        reader = PyPDF2.PdfReader(pdf_file)
-        text = ""
-        for page in reader.pages:
-            page_text = page.extract_text() or ""
-            text += page_text + "\n"
-        logger.info(f"Extracted text length: {len(text)} characters")
-        logger.debug(f"Extracted text sample: {text[:500]}")
-        if not text.strip():
-            return "Error: No text extracted from PDF."
-        return text
-    except Exception as e:
-        logger.error(f"Text extraction error: {str(e)}")
-        return f"Error extracting text: {str(e)}"
-
-def parse_contract(text):
-    """Parse contract text into clauses and classify risks using keyword-based heuristic."""
-    # Clean text: replace multiple newlines with single, handle LaTeX artifacts
-    text = text.replace("\n\n", "\n").replace("\t", " ")
-    sentences = nltk.sent_tokenize(text)
-    logger.info(f"Number of sentences tokenized: {len(sentences)}")
-    logger.debug(f"Sample sentences: {sentences[:3]}")
-
+    reader = PdfReader(pdf_file)
+    text = ""
+    for page in reader.pages:
+        text += page.extract_text() + "\n"
+    return text
+
+# Split text into clauses
+def split_into_clauses(text):
+    clauses = re.split(r'\n\s*\d+\.\s*|\n\s*[A-Z]\.\s*', text)
+    clauses = [clause.strip() for clause in clauses if clause.strip()]
+    return clauses
+
+# Load BERT model and tokenizer
+@st.cache_resource
+def load_model():
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=3)  # Fine-tuned for 3 risk levels
+    return tokenizer, model
+
+# Process clauses and assign risk scores
+def process_clauses(clauses, tokenizer, model):
     results = []
-    risk_scores = []
-
-    for idx, sentence in enumerate(sentences):
-        sentence = sentence.strip()
-        if len(sentence) < 10:  # Skip short sentences
-            logger.debug(f"Skipping short sentence (length {len(sentence)}): {sentence}")
-            continue
-
-        # Heuristic classification based on keywords
-        sentence_lower = sentence.lower()
-        clause_type = None
-        for c_type, keywords in KEYWORD_MAP.items():
-            if any(keyword in sentence_lower for keyword in keywords):
-                clause_type = c_type
-                break
-
-        if clause_type not in CLAUSE_TYPES:
-            logger.debug(f"No relevant clause type for sentence {idx}: {sentence}")
-            continue
-
-        # Assign a dummy score based on keyword presence (simulating model confidence)
-        score = RISK_WEIGHTS[clause_type] * 0.9  # 0.9 as a dummy confidence score
+    risk_levels = {0: 'Low', 1: 'Medium', 2: 'High'}
+    
+    for clause in clauses:
+        inputs = tokenizer(clause, return_tensors="pt", truncation=True, padding=True, max_length=512)
+        with torch.no_grad():
+            outputs = model(**inputs)
+        logits = outputs.logits
+        risk_score = torch.softmax(logits, dim=1).numpy()[0]
+        risk_level = risk_levels[np.argmax(risk_score)]
+        
         results.append({
-            "clause_id": idx,
-            "text": sentence,
-            "clause_type": clause_type,
-            "risk_score": round(score, 2)
+            'clause_text': clause,
+            'risk_level': risk_level,
+            'severity_score': float(np.max(risk_score)),
+            'clause_type': infer_clause_type(clause)  # Simplified clause type inference
+        })
+    
+    return results
+
+# Simplified clause type inference (extend with more sophisticated logic as needed)
+def infer_clause_type(clause):
+    if 'liability' in clause.lower():
+        return 'Liability'
+    elif 'payment' in clause.lower():
+        return 'Payment'
+    else:
+        return 'General'
+
+# Save results to Salesforce
+def save_to_salesforce(sf, results, contract_id):
+    for result in results:
+        sf.Contract_Risk__c.create({
+            'Contract__c': contract_id,
+            'Clause_Text__c': result['clause_text'][:255],  # Truncate if needed
+            'Risk_Level__c': result['risk_level'],
+            'Severity_Score__c': result['severity_score'],
+            'Clause_Type__c': result['clause_type']
         })
-        risk_scores.append(score)
-        logger.info(f"Detected clause {idx}: {clause_type} with risk score {score}")
-
-    return results, risk_scores
 
-def generate_heatmap(risk_scores):
-    """Generate heatmap for risk scores."""
-    if not risk_scores:
-        logger.warning("No risk scores to generate heatmap.")
-        return None
-    data = np.array(risk_scores).reshape(1, -1)
+# Generate heatmap
+def generate_heatmap(results):
+    df = pd.DataFrame(results)
+    risk_scores = df['severity_score'].values
     plt.figure(figsize=(10, 2))
-    sns.heatmap(data, cmap="YlOrRd", annot=True, fmt=".2f", cbar_kws={'label': 'Risk Score'})
-    plt.title("Contract Risk Heatmap")
-    plt.xlabel("Clause Index")
-    plt.ylabel("Risk")
-    buffer = BytesIO()
-    plt.savefig(buffer, format="png", bbox_inches="tight")
-    plt.close()
-    buffer.seek(0)
-    return buffer
-
-def generate_pdf_report(results, heatmap_buffer):
-    """Generate PDF report with summary and heatmap."""
-    buffer = BytesIO()
-    c = canvas.Canvas(buffer, pagesize=letter)
-    c.setFont("Helvetica", 12)
-    c.drawString(50, 750, "Contract Risk Analysis Report")
+    sns.heatmap([risk_scores], cmap='RdYlGn_r', annot=True, fmt='.2f', cbar_kws={'label': 'Risk Severity'})
+    plt.title('Contract Clause Risk Heatmap')
+    plt.xlabel('Clause Index')
+    plt.yticks([])
+    st.pyplot(plt)
+
+# Streamlit interface
+def main():
+    st.title("Contract Risk Analyzer")
     
-    # Summary
-    c.drawString(50, 720, "Summary of Risk-Prone Clauses:")
-    y = 700
-    for result in results[:5]:  # Limit to top 5 for brevity
-        text = f"Clause {result['clause_id']}: {result['clause_type'].capitalize()} (Risk: {result['risk_score']})"
-        c.drawString(50, y, text[:80] + "..." if len(text) > 80 else text)
-        y -= 20
-
-    # Embed heatmap
-    if heatmap_buffer:
-        c.drawImage(BytesIO(heatmap_buffer.read()), 50, y-200, width=500, height=100)
+    # File upload
+    uploaded_file = st.file_uploader("Upload Contract PDF", type=["pdf"])
+    contract_id = st.text_input("Enter Contract ID")
     
-    c.showPage()
-    c.save()
-    buffer.seek(0)
-    return buffer
-
-def process_contract(pdf_file):
-    """Main function to process uploaded contract."""
-    # Extract text
-    text = extract_text_from_pdf(pdf_file)
-    if "Error" in text:
-        return text, None, None, {"Error": text}
-
-    # Parse and classify
-    results, risk_scores = parse_contract(text)
-    if not results:
-        return "No relevant clauses detected.", None, None, {"Summary": "No risk-prone clauses found."}
-
-    # Generate outputs
-    json_output = json.dumps(results, indent=2)
-    heatmap_buffer = generate_heatmap(risk_scores)
-    pdf_report = generate_pdf_report(results, heatmap_buffer)
-
-    return json_output, heatmap_buffer, pdf_report, {"Summary": f"Detected {len(results)} risk-prone clauses."}
-
-# Gradio interface
-iface = gr.Interface(
-    fn=process_contract,
-    inputs=gr.File(label="Upload Contract PDF"),
-    outputs=[
-        gr.Textbox(label="JSON Output"),
-        gr.Image(label="Risk Heatmap"),
-        gr.File(label="Download PDF Report"),
-        gr.JSON(label="Summary")
-    ],
-    title="Contract Risk Analyzer",
-    description="Upload a contract PDF to analyze risk-prone clauses and visualize results."
-)
+    if uploaded_file and contract_id:
+        # Extract and process text
+        text = extract_text_from_pdf(uploaded_file)
+        clauses = split_into_clauses(text)
+        
+        # Load model and process clauses
+        tokenizer, model = load_model()
+        results = process_clauses(clauses, tokenizer, model)
+        
+        # Display results
+        st.subheader("Clause Analysis Results")
+        for i, result in enumerate(results, 1):
+            st.write(f"**Clause {i}**")
+            st.write(f"Text: {result['clause_text'][:100]}...")
+            st.write(f"Clause Type: {result['clause_type']}")
+            st.write(f"Risk Level: {result['risk_level']}")
+            st.write(f"Severity Score: {result['severity_score']:.2f}")
+            st.write("---")
+        
+        # Generate and display heatmap
+        generate_heatmap(results)
+        
+        # Save to Salesforce
+        if st.button("Save to Salesforce"):
+            sf = connect_to_salesforce()
+            save_to_salesforce(sf, results, contract_id)
+            st.success("Results saved to Salesforce!")
 
 if __name__ == "__main__":
-    iface.launch()
+    main()