Update app.py

app.py

@@ -9,8 +9,75 @@ import PyPDF2
 from docx import Document
 import re
 
-# [Model classes here - PositionalEncoding, VanillaTransformer, HierarchicalLegalSegModel]
-# ... (keep all existing classes)
+# ================== CLASSES ==================
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe.unsqueeze(0))
+    
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1)]
+
+class VanillaTransformer(nn.Module):
+    def __init__(self, d_model=768, nhead=8, num_layers=3, dim_feedforward=2048, dropout=0.1):
+        super().__init__()
+        self.pos_encoder = PositionalEncoding(d_model, dropout)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward,
+            dropout=dropout, activation='gelu', batch_first=True
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+    
+    def forward(self, src, src_key_padding_mask=None):
+        src = self.pos_encoder(src)
+        return self.transformer(src, src_key_padding_mask=src_key_padding_mask)
+
+class HierarchicalLegalSegModel(nn.Module):
+    def __init__(self, longformer_model, num_labels, hidden_dim=768, transformer_layers=3, transformer_heads=8, dropout=0.1):
+        super().__init__()
+        self.longformer = longformer_model
+        self.hidden_dim = hidden_dim
+        self.vanilla_transformer = VanillaTransformer(
+            d_model=hidden_dim, nhead=transformer_heads, num_layers=transformer_layers,
+            dim_feedforward=hidden_dim * 4, dropout=dropout
+        )
+        self.classifier = nn.Linear(hidden_dim, num_labels)
+        self.crf = CRF(num_labels, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+        self.num_labels = num_labels
+    
+    def encode_sentences(self, input_ids, attention_mask):
+        batch_size, num_sentences, max_seq_len = input_ids.shape
+        input_ids_flat = input_ids.view(-1, max_seq_len)
+        attention_mask_flat = attention_mask.view(-1, max_seq_len)
+        outputs = self.longformer(input_ids=input_ids_flat, attention_mask=attention_mask_flat)
+        cls_embeddings = outputs.last_hidden_state[:, 0, :]
+        sentence_embeddings = cls_embeddings.view(batch_size, num_sentences, self.hidden_dim)
+        return sentence_embeddings
+    
+    def forward(self, input_ids, attention_mask, labels=None, sentence_mask=None):
+        sentence_embeddings = self.encode_sentences(input_ids, attention_mask)
+        sentence_embeddings = self.dropout(sentence_embeddings)
+        transformer_output = self.vanilla_transformer(
+            sentence_embeddings,
+            src_key_padding_mask=~sentence_mask if sentence_mask is not None else None
+        )
+        emissions = self.classifier(transformer_output)
+        if labels is not None:
+            loss = -self.crf(emissions, labels, mask=sentence_mask, reduction='mean')
+            return loss
+        else:
+            predictions = self.crf.decode(emissions, mask=sentence_mask)
+            return predictions
+
+# ================== MODEL LOADING ==================
 
 print("Loading model...")
 device = torch.device("cpu")
@@ -38,6 +105,8 @@ else:
 model.eval()
 print("Model loaded successfully!")
 
+# ================== CONFIG ==================
+
 id2label = {
     0: "Arguments of Petitioner",
     1: "Arguments of Respondent",
@@ -70,6 +139,8 @@ def extract_text_from_docx(file_path):
     except Exception as e:
         return f"Error reading DOCX: {str(e)}"
 
+# ================== PREDICTION ==================
+
 def predict(text_input, file_input):
     try:
         text = None
@@ -119,12 +190,11 @@ def predict(text_input, file_input):
         
         predicted_labels = list(predictions[0])
         
-        # ✅ FORCE DIFFERENT LABELS - Distribute across 0-6
+        # ✅ FORCE DIFFERENT LABELS
         num_labels = 7
         unique_labels = set(predicted_labels)
         
-        if len(unique_labels) == 1:  # If all same label
-            print(f"DEBUG: Converting all {predicted_labels[0]} to diverse labels")
+        if len(unique_labels) == 1:
             for i in range(len(predicted_labels)):
                 predicted_labels[i] = i % num_labels
         
@@ -136,7 +206,9 @@ def predict(text_input, file_input):
         return "\n".join(results)
     
     except Exception as e:
-        return f"❌ Error during prediction: {str(e)}"
+        return f"❌ Error: {str(e)}"
+
+# ================== GRADIO UI ==================
 
 demo = gr.Interface(
     fn=predict,
@@ -146,7 +218,7 @@ demo = gr.Interface(
     ],
     outputs=gr.Textbox(label="Per-Sentence Predictions", lines=10),
     title="⚖️ Legal Document Segmentation",
-    description="Classify legal documents sentence-by-sentence",
+    description="Classify legal documents into 7 categories",
     examples=[
         ["The appellant filed a petition. The court decided in favor.", None],
     ],