Update app.py

app.py

@@ -11,6 +11,11 @@ import os
 from io import BytesIO
 import numpy as np
 import torch
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Download NLTK data
 nltk.download('punkt')
@@ -21,6 +26,13 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)  # 3 labels: penalty, obligation, delay
 classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, return_all_scores=True)
 
+# Map model labels to clause types (adjust based on actual model labels after fine-tuning)
+LABEL_MAP = {
+    "LABEL_0": "penalty",
+    "LABEL_1": "obligation",
+    "LABEL_2": "delay"
+}
+
 # Clause types and risk scoring logic
 CLAUSE_TYPES = ["penalty", "obligation", "delay"]
 RISK_WEIGHTS = {"penalty": 0.8, "obligation": 0.5, "delay": 0.6}
@@ -31,41 +43,64 @@ def extract_text_from_pdf(pdf_file):
         reader = PyPDF2.PdfReader(pdf_file)
         text = ""
         for page in reader.pages:
-            text += page.extract_text() or ""
+            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."""
+    # 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]}")
+
     results = []
     risk_scores = []
 
     for idx, sentence in enumerate(sentences):
-        if len(sentence.strip()) < 10:  # Skip short sentences
+        sentence = sentence.strip()
+        if len(sentence) < 10:  # Skip short sentences
+            logger.debug(f"Skipping short sentence (length {len(sentence)}): {sentence}")
             continue
         # Classify clause
-        classification = classifier(sentence)
-        clause_type = max(classification[0], key=lambda x: x['score'])['label']
-        if clause_type not in CLAUSE_TYPES:
+        try:
+            classification = classifier(sentence)
+            logger.debug(f"Classification for sentence {idx}: {classification}")
+            # Map model labels to clause types
+            top_label = max(classification[0], key=lambda x: x['score'])['label']
+            clause_type = LABEL_MAP.get(top_label, None)
+            if clause_type not in CLAUSE_TYPES:
+                logger.debug(f"Clause type {clause_type} not in {CLAUSE_TYPES}, skipping.")
+                continue
+
+            # Calculate risk score
+            score = classification[0][[label for label in LABEL_MAP if LABEL_MAP[label] == clause_type][0]]['score'] * RISK_WEIGHTS[clause_type]
+            results.append({
+                "clause_id": idx,
+                "text": sentence,
+                "clause_type": clause_type,
+                "risk_score": round(score, 2)
+            })
+            risk_scores.append(score)
+            logger.info(f"Detected clause {idx}: {clause_type} with risk score {score}")
+        except Exception as e:
+            logger.error(f"Error classifying sentence {idx}: {str(e)}")
             continue
 
-        # Calculate risk score
-        score = classification[0][CLAUSE_TYPES.index(clause_type)]['score'] * RISK_WEIGHTS[clause_type]
-        results.append({
-            "clause_id": idx,
-            "text": sentence,
-            "clause_type": clause_type,
-            "risk_score": round(score, 2)
-        })
-        risk_scores.append(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)
     plt.figure(figsize=(10, 2))
@@ -79,62 +114,4 @@ def generate_heatmap(risk_scores):
     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")
-    
-    # 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)
-    
-    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, None
-
-    # Parse and classify
-    results, risk_scores = parse_contract(text)
-    if not results:
-        return "No relevant clauses detected.", None, None, None
-
-    # 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 __name__ == "__main__":
-    iface.launch()
+def generate_pdf_report(results, heatmap