Spaces:

nlst
/

bert_testing_1

Paused

App Files Files Community

faheem66 commited on Jul 30, 2024

Commit

11402c8

1 Parent(s): d2d9304

added the bert model with synthetic data for initial training and testing

Browse files

Files changed (2) hide show

app.py +223 -0
requirements.txt +6 -0

app.py ADDED Viewed

	@@ -0,0 +1,223 @@

+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import BertTokenizer, BertForSequenceClassification, AdamW
+from sklearn.model_selection import train_test_split
+import gradio as gr
+import random
+from faker import Faker
+import html
+import numpy as np
+from tqdm import tqdm
+# Constants
+MAX_LENGTH = 512
+BATCH_SIZE = 16
+EPOCHS = 5
+LEARNING_RATE = 2e-5
+fake = Faker()
+def generate_employee():
+    name = fake.name()
+    job = fake.job()
+    ext = f"ext. {random.randint(1000, 9999)}"
+    email = f"{name.lower().replace(' ', '.')}@example.com"
+    return name, job, ext, email
+def generate_html_content(num_employees=9):
+    employees = [generate_employee() for _ in range(num_employees)]
+    html_content = f"""
+    <html>
+    <head>
+        <title>Employee Directory</title>
+    </head>
+    <body>
+        <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">'
+        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>
+            </div>
+        """
+        if (i + 1) % 3 == 0 or i == len(employees) - 1:
+            html_content += '</div>'
+    html_content += """
+        </div>
+    </body>
+    </html>
+    """
+    return html_content
+def generate_dataset(num_samples=1000):
+    dataset = []
+    for _ in range(num_samples):
+        html_content = generate_html_content()
+        employees = []
+        for line in html_content.split('\n'):
+            if '<strong>' in line:
+                name = line.split('<strong>')[1].split('</strong>')[0]
+            elif '<span style="font-size: 16px">' in line:
+                if 'ext.' in line:
+                    ext = line.split('<span style="font-size: 16px">')[1].split('</span>')[0]
+                else:
+                    job = line.split('<span style="font-size: 16px">')[1].split('</span>')[0]
+            elif '<a href="mailto:' in line:
+                email = line.split('<a href="mailto:')[1].split('">')[0]
+                employees.append(f"{name}\n{job}\n{ext}\n{email}")
+        dataset.append((html_content, '\n\n'.join(employees)))
+    return dataset
+class HTMLDataset(Dataset):
+    def __init__(self, data, tokenizer, max_length):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+    def __len__(self):
+        return len(self.data)
+    def __getitem__(self, idx):
+        html, extracted = self.data[idx]
+        encoding = self.tokenizer.encode_plus(
+            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',
+        )
+        return {
+            'input_ids': encoding['input_ids'].flatten(),
+            'attention_mask': encoding['attention_mask'].flatten(),
+            'labels': torch.tensor(extracted, dtype=torch.float)
+        }
+def train_model(progress=gr.Progress()):
+    # Generate synthetic dataset
+    dataset = generate_dataset(num_samples=1000)
+    train_data, val_data = train_test_split(dataset, test_size=0.2, random_state=42)
+    # Initialize tokenizer and model
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=1)
+    # Prepare datasets and dataloaders
+    train_dataset = HTMLDataset(train_data, tokenizer, MAX_LENGTH)
+    val_dataset = HTMLDataset(val_data, tokenizer, MAX_LENGTH)
+    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
+    val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE)
+    # Initialize optimizer
+    optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)
+    # Training loop
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model.to(device)
+    for epoch in progress.tqdm(range(EPOCHS), desc="Training Progress"):
+        model.train()
+        train_loss = 0
+        for batch in train_dataloader:
+            input_ids = batch['input_ids'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            labels = batch['labels'].to(device)
+            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
+            loss = outputs.loss
+            train_loss += loss.item()
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+        # Validation
+        model.eval()
+        val_loss = 0
+        with torch.no_grad():
+            for batch in val_dataloader:
+                input_ids = batch['input_ids'].to(device)
+                attention_mask = batch['attention_mask'].to(device)
+                labels = batch['labels'].to(device)
+                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
+                val_loss += outputs.loss.item()
+        avg_train_loss = train_loss / len(train_dataloader)
+        avg_val_loss = val_loss / len(val_dataloader)
+        progress(f"Epoch {epoch + 1}/{EPOCHS}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")
+    return model, tokenizer
+def extract_content(html, model, tokenizer):
+    model.eval()
+    encoding = tokenizer.encode_plus(
+        html,
+        add_special_tokens=True,
+        max_length=MAX_LENGTH,
+        return_token_type_ids=False,
+        padding='max_length',
+        truncation=True,
+        return_attention_mask=True,
+        return_tensors='pt',
+    )
+    input_ids = encoding['input_ids'].to(model.device)
+    attention_mask = encoding['attention_mask'].to(model.device)
+    with torch.no_grad():
+        outputs = model(input_ids, attention_mask=attention_mask)
+        predictions = outputs.logits.sigmoid().cpu().numpy()
+    # Extract content based on predictions
+    # This is a placeholder implementation and needs to be adjusted based on your specific use case
+    extracted_content = f"Extracted content (confidence: {predictions[0][0]:.2f})"
+    return extracted_content
+def gradio_interface(html_input):
+    global trained_model, trained_tokenizer
+    extracted_content = extract_content(html_input, trained_model, trained_tokenizer)
+    return extracted_content
+print("Starting training process...")
+trained_model, trained_tokenizer = train_model()
+print("Training completed. Launching Gradio interface...")
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Textbox(lines=10, label="Input HTML"),
+    outputs=gr.Textbox(label="Extracted Content"),
+    title="HTML Content Extractor",
+    description="Enter HTML content to extract information."
+)
+iface.launch()

requirements.txt ADDED Viewed

	@@ -0,0 +1,6 @@

+torch
+transformers
+scikit-learn
+faker
+gradio
+tqdm