changed training approach

app.py

@@ -1,13 +1,15 @@
 import torch
 from torch.utils.data import Dataset, DataLoader
-from transformers import BertTokenizer, BertForSequenceClassification, AdamW
+from transformers import BertTokenizer, BertForTokenClassification, AdamW
 from sklearn.model_selection import train_test_split
 import gradio as gr
 import random
 from faker import Faker
 import html
+import json
 import numpy as np
 from tqdm import tqdm
+import os
 
 # Constants
 MAX_LENGTH = 512
@@ -26,7 +28,7 @@ def generate_employee():
     return name, job, ext, email
 
 
-def generate_html_content(num_employees=9):
+def generate_html_content(num_employees=3):
     employees = [generate_employee() for _ in range(num_employees)]
 
     html_content = f"""
@@ -38,26 +40,22 @@ def generate_html_content(num_employees=9):
         <div class="row ts-three-column-row standard-row">
     """
 
-    for i, (name, job, ext, email) in enumerate(employees):
-        if i % 3 == 0:
-            html_content += '<div class="column ts-three-column">'
-
+    for name, job, ext, email in employees:
         html_content += f"""
-            <div class="block">
-                <div class="text-block" style="text-align: center;">
-                    <p>
-                        <strong>{html.escape(name)}</strong><br>
-                        <span style="font-size: 16px">{html.escape(job)}</span><br>
-                        <span style="font-size: 16px">{html.escape(ext)}</span><br>
-                        <a href="mailto:{html.escape(email)}">Send Email</a>
-                    </p>
+            <div class="column ts-three-column">
+                <div class="block">
+                    <div class="text-block" style="text-align: center;">
+                        <p>
+                            <strong>{html.escape(name)}</strong><br>
+                            <span style="font-size: 16px">{html.escape(job)}</span><br>
+                            <span style="font-size: 16px">{html.escape(ext)}</span><br>
+                            <a href="mailto:{html.escape(email)}">Send Email</a>
+                        </p>
+                    </div>
                 </div>
             </div>
         """
 
-        if (i + 1) % 3 == 0 or i == len(employees) - 1:
-            html_content += '</div>'
-
     html_content += """
         </div>
     </body>
@@ -80,6 +78,8 @@ class HTMLDataset(Dataset):
         self.data = data
         self.tokenizer = tokenizer
         self.max_length = max_length
+        self.label_map = {"O": 0, "B-NAME": 1, "I-NAME": 2, "B-JOB": 3, "I-JOB": 4, "B-EXT": 5, "I-EXT": 6,
+                          "B-EMAIL": 7, "I-EMAIL": 8}
 
     def __len__(self):
         return len(self.data)
@@ -90,22 +90,38 @@ class HTMLDataset(Dataset):
             html,
             add_special_tokens=True,
             max_length=self.max_length,
-            return_token_type_ids=False,
             padding='max_length',
             truncation=True,
             return_attention_mask=True,
             return_tensors='pt',
         )
 
-        # Create a binary classification task: 1 if there are employees, 0 otherwise
-        label = 1 if employees else 0
+        labels = self.create_labels(encoding['input_ids'][0], employees)
 
         return {
             'input_ids': encoding['input_ids'].flatten(),
             'attention_mask': encoding['attention_mask'].flatten(),
-            'labels': torch.tensor(label, dtype=torch.float)
+            'labels': torch.tensor(labels, dtype=torch.long)
         }
 
+    def create_labels(self, tokens, employees):
+        labels = [0] * len(tokens)  # Initialize all labels as "O"
+        for name, job, ext, email in employees:
+            self.label_sequence(tokens, name, "NAME", labels)
+            self.label_sequence(tokens, job, "JOB", labels)
+            self.label_sequence(tokens, ext, "EXT", labels)
+            self.label_sequence(tokens, email, "EMAIL", labels)
+        return labels
+
+    def label_sequence(self, tokens, text, label_type, labels):
+        text_tokens = self.tokenizer.encode(text, add_special_tokens=False)
+        for i in range(len(tokens) - len(text_tokens) + 1):
+            if tokens[i:i + len(text_tokens)] == text_tokens:
+                labels[i] = self.label_map[f"B-{label_type}"]
+                for j in range(1, len(text_tokens)):
+                    labels[i + j] = self.label_map[f"I-{label_type}"]
+                break
+
 
 def train_model():
     # Generate synthetic dataset
@@ -114,7 +130,7 @@ def train_model():
 
     # Initialize tokenizer and model
     tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-    model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=1)
+    model = BertForTokenClassification.from_pretrained('bert-base-uncased', num_labels=9)
 
     # Prepare datasets and dataloaders
     train_dataset = HTMLDataset(train_data, tokenizer, MAX_LENGTH)
@@ -171,7 +187,6 @@ def extract_content(html, model, tokenizer):
         html,
         add_special_tokens=True,
         max_length=MAX_LENGTH,
-        return_token_type_ids=False,
         padding='max_length',
         truncation=True,
         return_attention_mask=True,
@@ -183,16 +198,50 @@ def extract_content(html, model, tokenizer):
 
     with torch.no_grad():
         outputs = model(input_ids, attention_mask=attention_mask)
-        predictions = outputs.logits.sigmoid().cpu().numpy()
-
-    # Extract content based on predictions
-    confidence = predictions[0][0]
-    if confidence > 0.5:
-        extracted_content = "Employee information detected"
-    else:
-        extracted_content = "No employee information detected"
-
-    return f"{extracted_content} (confidence: {confidence:.2f})"
+        predictions = outputs.logits.argmax(dim=2)
+
+    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
+    label_map = {0: "O", 1: "B-NAME", 2: "I-NAME", 3: "B-JOB", 4: "I-JOB", 5: "B-EXT", 6: "I-EXT", 7: "B-EMAIL",
+                 8: "I-EMAIL"}
+
+    extracted_info = []
+    current_entity = {"type": None, "value": ""}
+
+    for token, prediction in zip(tokens, predictions[0]):
+        if token == "[PAD]":
+            break
+
+        label = label_map[prediction.item()]
+        if label.startswith("B-"):
+            if current_entity["type"]:
+                extracted_info.append(current_entity)
+            current_entity = {"type": label[2:], "value": token}
+        elif label.startswith("I-"):
+            if current_entity["type"] == label[2:]:
+                current_entity["value"] += " " + token
+        elif label == "O":
+            if current_entity["type"]:
+                extracted_info.append(current_entity)
+                current_entity = {"type": None, "value": ""}
+
+    if current_entity["type"]:
+        extracted_info.append(current_entity)
+
+    # Group entities by employee
+    employees = []
+    current_employee = {}
+    for entity in extracted_info:
+        if entity["type"] == "NAME":
+            if current_employee:
+                employees.append(current_employee)
+            current_employee = {"name": entity["value"]}
+        else:
+            current_employee[entity["type"].lower()] = entity["value"]
+
+    if current_employee:
+        employees.append(current_employee)
+
+    return json.dumps(employees, indent=2)
 
 
 def test_model(html_input, model, tokenizer):
@@ -201,18 +250,30 @@ def test_model(html_input, model, tokenizer):
 
 
 def gradio_interface(html_input, test_type):
-    global trained_model, trained_tokenizer
+    global model, tokenizer
     if test_type == "Custom Input":
-        return test_model(html_input, trained_model, trained_tokenizer)
+        result = test_model(html_input, model, tokenizer)
+        return html_input, result
     elif test_type == "Generate Random HTML":
         random_html, _ = generate_html_content()
-        result = test_model(random_html, trained_model, trained_tokenizer)
+        result = test_model(random_html, model, tokenizer)
         return random_html, result
 
 
-print("Starting training process...")
-trained_model, trained_tokenizer = train_model()
-print("Training completed. Launching Gradio interface...")
+# Check if the model is already trained and saved
+if os.path.exists('model.pth') and os.path.exists('tokenizer'):
+    print("Loading pre-trained model...")
+    model = BertForTokenClassification.from_pretrained('bert-base-uncased', num_labels=9)
+    model.load_state_dict(torch.load('model.pth'))
+    tokenizer = BertTokenizer.from_pretrained('tokenizer')
+else:
+    print("Training new model...")
+    model, tokenizer = train_model()
+    # Save the model and tokenizer
+    torch.save(model.state_dict(), 'model.pth')
+    tokenizer.save_pretrained('tokenizer')
+
+print("Launching Gradio interface...")
 
 iface = gr.Interface(
     fn=gradio_interface,
@@ -222,10 +283,10 @@ iface = gr.Interface(
     ],
     outputs=[
         gr.Textbox(lines=10, label="HTML Content"),
-        gr.Textbox(label="Extracted Information")
+        gr.Textbox(label="Extracted Information (JSON)")
     ],
     title="HTML Content Extractor",
-    description="Enter HTML content or generate random HTML to test the model."
+    description="Enter HTML content or generate random HTML to test the model. The model will extract employee information and return it in JSON format."
 )
 
-iface.launch()
+iface.launch(share=True)