Build a Form Field Classifier with PyTorch
Train a neural network from scratch to classify job application form fields.
Introduction
30 Second Summary
Job applications are tedious. You fill out the same information over and over across dozens of forms, each with slightly different field names and layouts.
In this project, you will build a custom neural network from scratch that learns to classify job application form fields from raw HTML metadata. Your trained model will power an intelligent autofill system that understands what each form field expects.
What You'll Build
A complete ML pipeline that takes form field metadata (label text, input types, name attributes) and predicts what category of information each field expects.
By the end of this project, you'll have:
- A custom PyTorch neural network trained on your GPU to classify 18 form field categories with 85%+ accuracy.
- A synthetic training dataset of 2000+ labeled samples representing real job application form patterns.
- A running FastAPI prediction endpoint that accepts form field metadata and returns predicted categories, ready for a future Chrome extension to call.
- A trained model registered in the DataEngineX model registry with versioning and stage management.
- Secret Mission: Fine-tune DistilBERT for semantic understanding of ambiguous form fields that the rule-based classifier misses.
Are there any prerequisites?
You need Python 3.13 installed, an NVIDIA GPU with drivers configured, and familiarity with Python classes and functions.
This project uses an NVIDIA Quadro T2000 (4GB VRAM). Any NVIDIA GPU with Turing architecture or newer will work.
Before We Start
You are building a custom neural network that understands job application forms. Given raw metadata about a form field (its label, input type, name attribute, and placeholder text), your model predicts what kind of information that field expects: first name, email, phone number, work experience, cover letter, and more.
Before you can train anything, you need the right tools in place. In this step, you will install PyTorch with GPU acceleration, scaffold your project directory, and wire up DataEngineX for ML tracking.
In this step, get ready to:
- Install PyTorch with CUDA 12.6 GPU support and verify your Quadro T2000 is accessible.
- Scaffold the project directory structure and install dependencies.
- Configure DataEngineX and verify the engine loads.
Install PyTorch with CUDA 12.6
Your NVIDIA Quadro T2000 uses the Turing architecture (compute capability 7.5). PyTorch's CUDA 12.6 build supports Turing GPUs, so you need to install from that specific index.
- Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type Terminal and press Enter to open it.
- Install PyTorch with CUDA 12.6 support by running this command:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126What does this command do?
The --index-url flag tells pip to pull packages from PyTorch's CUDA 12.6 wheel repository instead of the default CPU-only builds on PyPI.
This ensures your PyTorch installation can offload tensor operations to your GPU, which makes training dramatically faster than CPU-only execution.
- The installation downloads several hundred megabytes. Wait for it to complete before continuing.
- Verify your GPU is accessible by running this command:
python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"✔️ I see True and Quadro T2000
Your GPU is ready for training. PyTorch can see your Quadro T2000 and will use it to accelerate tensor computations.
ⓧ I see False
PyTorch installed but cannot detect your GPU. This usually means either your NVIDIA drivers are outdated or the CPU-only build was installed instead of the CUDA build.
- Confirm your NVIDIA drivers are installed by running:
nvidia-smi- If that shows your GPU, uninstall and reinstall PyTorch to ensure the CUDA build:
pip uninstall torch torchvision torchaudio -y
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126- If nvidia-smi is not found, you need to install NVIDIA GPU drivers for your Quadro T2000 before continuing.
ⓧ ModuleNotFoundError
The installation did not complete successfully. Check your terminal output for errors.
- Make sure you are using Python 3.13 by running:
python --version- Re-run the install command and watch for any error messages about disk space or network connectivity.
Scaffold the project structure
- Move to your Desktop by running this command:
cd ~/Desktop- Create the project folder and subdirectories by running these commands:
mkdir formfield-classifier
cd formfield-classifier
mkdir -p src/formfield_classifier data models- Confirm the directories exist by running:
ls- You should see data, models, and src in the output.
Don't see all three folders?
Make sure you ran the commands from inside the formfield-classifier folder. Run pwd to check your current location.
The -p flag in mkdir -p creates parent directories automatically. If src/formfield_classifier is missing, re-run the mkdir command.
Help me debug my project folder structure.
- Install the project dependencies by running:
pip install dataenginex fastapi uvicorn[standard]What are these packages?
- dataenginex is the DataEngineX engine that handles ML tracking, model registry, and pipeline orchestration.
- fastapi is a modern Python web framework you will use to serve predictions as a REST API.
- uvicorn[standard] is the ASGI server that runs your FastAPI app.
Seeing installation errors?
Make sure you are running Python 3.13. DataEngineX requires Python 3.13 or higher.
If you see a permissions error, try adding --user to the end of the pip install command.
Help me fix my pip install errors.
- Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type Visual Studio Code and press Enter to open it.
- Click File in the top menu bar.
- Click Open Folder.
- Navigate to Desktop/formfield-classifier and click Open.
- In the Explorer sidebar, right-click on the src/formfield_classifier folder.
- Click New File and name it __init__.py.
- Add this code to the file:
"""Form Field Classifier - PyTorch model for job application form field classification."""
__version__ = "0.1.0"- Save __init__.py with Cmd+S (macOS) or Ctrl+S (Windows).
Create the DataEngineX configuration
DataEngineX uses a dex.yaml file at the project root to know how to store training metrics and model artifacts. This config tells the engine where to persist experiment data on your local machine.
- In VS Code's Explorer sidebar, right-click on the empty area below your folders.
- Click New File and name it dex.yaml.
- Add this configuration to the file:
project:
name: formfield-classifier
version: "0.1.0"
description: "Form field classification model for job application autofill"
storage:
backend: local
path: ./data
ml:
tracking:
backend: json
path: ./.dex/experimentsWhat does this config do?
- project gives your project a name and version that DataEngineX uses for tracking.
- storage tells the engine to use local file storage in your ./data directory.
- ml.tracking configures where training experiment logs are saved. The engine persists run data to .dex/store.duckdb automatically.
- Save dex.yaml.
- Open a terminal inside VS Code by pressing Ctrl+` (backtick key, same on macOS and Windows).
- Verify the DataEngineX engine loads by running this command:
python -c "from dataenginex.engine import DexEngine; engine = DexEngine('dex.yaml'); print('Engine loaded:', engine)"- You should see Engine loaded: followed by the engine object representation. No errors means DataEngineX is configured correctly.
Seeing an import error or config error?
If you see ModuleNotFoundError, make sure the pip install dataenginex command completed successfully earlier.
If the engine fails to load, double check that dex.yaml is saved at the project root (not inside a subfolder). Your terminal's working directory should be formfield-classifier/.
Help me fix my DexEngine loading error.
✔️ Awesome, I've got everything!
Great. Double check you have saved both files before moving on.
ⓧ I'd like to double check the full code
Here is the full dex.yaml for comparison:
project:
name: formfield-classifier
version: "0.1.0"
description: "Form field classification model for job application autofill"
storage:
backend: local
path: ./data
ml:
tracking:
backend: json
path: ./.dex/experimentsAnd the full src/formfield_classifier/__init__.py:
"""Form Field Classifier - PyTorch model for job application form field classification."""
__version__ = "0.1.0"Your environment is locked and loaded. Next up, you will build a synthetic training dataset that teaches your neural network what job application form fields look like.
Create the Training Dataset
Your environment is locked and loaded. PyTorch can see your GPU, and DataEngineX loads your project config without errors.
Now you need data to train on. A neural network learns patterns from labeled examples, so you need a dataset of form field metadata (labels, input types, name attributes, placeholders) paired with the correct category for each field. In this step, you'll build a synthetic data generator that produces thousands of realistic training samples, then encode them into numerical tensors the network can consume.
In this step, get ready to:
- Define 18 target field categories with realistic metadata templates.
- Generate 2000+ labeled training samples with an 80/20 train/test split.
- Build a vocabulary and encode text features into fixed-length numerical tensors.
Define the field categories and data templates
Every machine learning model needs a clear set of labels it can predict. For job application forms, you need to capture the full range of fields a candidate encounters. You also need variation in how those fields appear across different websites.
- In VS Code, create a new file at src/formfield_classifier/dataset.py by right-clicking the src/formfield_classifier folder in the Explorer sidebar and selecting New File.
- Paste the following data definitions into the file:
"""Training dataset generation and loading for form field classification."""
import json
import random
from pathlib import Path
FIELD_CATEGORIES = [
"first_name", "last_name", "email", "phone", "linkedin",
"portfolio_url", "address", "city", "state", "zip_code",
"country", "company", "job_title", "start_date", "end_date",
"education", "skills", "cover_letter",
]
FIELD_TEMPLATES = {
"first_name": {
"labels": ["First Name", "First name", "Given Name", "Your first name", "Name (First)", "Forename"],
"input_types": ["text"],
"names": ["first_name", "firstName", "fname", "given_name", "first-name", "givenName"],
"placeholders": ["John", "Enter first name", "First name", "Your first name", ""],
},
"last_name": {
"labels": ["Last Name", "Last name", "Surname", "Family Name", "Name (Last)", "Your last name"],
"input_types": ["text"],
"names": ["last_name", "lastName", "lname", "surname", "family_name", "familyName"],
"placeholders": ["Doe", "Enter last name", "Last name", "Your surname", ""],
},
"email": {
"labels": ["Email", "Email Address", "E-mail", "Your Email", "Work Email", "Personal Email"],
"input_types": ["email", "text"],
"names": ["email", "emailAddress", "user_email", "applicant_email", "e-mail", "mail"],
"placeholders": ["john@example.com", "Enter email", "you@company.com", "Email address", ""],
},
"phone": {
"labels": ["Phone", "Phone Number", "Mobile", "Cell Phone", "Contact Number", "Telephone"],
"input_types": ["tel", "text"],
"names": ["phone", "phoneNumber", "mobile", "telephone", "cell", "phone_number"],
"placeholders": ["+1 (555) 123-4567", "Enter phone number", "Phone number", "Your phone", ""],
},
"linkedin": {
"labels": ["LinkedIn", "LinkedIn URL", "LinkedIn Profile", "Your LinkedIn", "LinkedIn Profile URL"],
"input_types": ["url", "text"],
"names": ["linkedin", "linkedinUrl", "linkedin_url", "linkedin_profile", "social_linkedin"],
"placeholders": ["https://linkedin.com/in/yourname", "LinkedIn URL", "Your LinkedIn profile", ""],
},
"portfolio_url": {
"labels": ["Portfolio", "Website", "Personal Website", "Portfolio URL", "Your Website", "GitHub"],
"input_types": ["url", "text"],
"names": ["portfolio", "website", "personal_url", "portfolio_url", "github", "url"],
"placeholders": ["https://yoursite.com", "Your website URL", "Portfolio link", ""],
},
"address": {
"labels": ["Address", "Street Address", "Home Address", "Address Line 1", "Street", "Mailing Address"],
"input_types": ["text"],
"names": ["address", "street_address", "address1", "street", "home_address", "addr"],
"placeholders": ["123 Main St", "Enter your address", "Street address", ""],
},
"city": {
"labels": ["City", "City/Town", "Your City", "Town", "Municipality"],
"input_types": ["text"],
"names": ["city", "town", "municipality", "city_name", "applicant_city"],
"placeholders": ["New York", "Enter city", "City", "Your city", ""],
},
"state": {
"labels": ["State", "State/Province", "Province", "Region", "State/Region"],
"input_types": ["text", "select"],
"names": ["state", "province", "region", "state_province", "applicant_state"],
"placeholders": ["California", "Select state", "State", "Province/State", ""],
},
"zip_code": {
"labels": ["Zip Code", "ZIP", "Postal Code", "Zip/Postal Code", "Postcode"],
"input_types": ["text"],
"names": ["zip", "zip_code", "postal_code", "zipCode", "postcode", "postalCode"],
"placeholders": ["10001", "Enter zip code", "Zip code", "Postal code", ""],
},
"country": {
"labels": ["Country", "Country/Region", "Your Country", "Nation", "Country of Residence"],
"input_types": ["select", "text"],
"names": ["country", "nation", "country_code", "residence_country", "applicant_country"],
"placeholders": ["United States", "Select country", "Country", ""],
},
"company": {
"labels": ["Company", "Company Name", "Employer", "Organization", "Current Company", "Previous Employer"],
"input_types": ["text"],
"names": ["company", "employer", "organization", "company_name", "org", "firm"],
"placeholders": ["Acme Corp", "Company name", "Your current employer", "Organization name", ""],
},
"job_title": {
"labels": ["Job Title", "Title", "Position", "Role", "Current Title", "Position Title"],
"input_types": ["text"],
"names": ["job_title", "title", "position", "role", "current_title", "jobTitle"],
"placeholders": ["Software Engineer", "Your job title", "Position title", "Current role", ""],
},
"start_date": {
"labels": ["Start Date", "From", "Date Started", "Begin Date", "Employment Start"],
"input_types": ["date", "text"],
"names": ["start_date", "startDate", "from_date", "date_started", "begin_date"],
"placeholders": ["MM/YYYY", "Start date", "When did you start?", "2020-01", ""],
},
"end_date": {
"labels": ["End Date", "To", "Date Ended", "End", "Employment End", "Currently Working"],
"input_types": ["date", "text"],
"names": ["end_date", "endDate", "to_date", "date_ended", "end", "until"],
"placeholders": ["MM/YYYY", "End date", "Present", "Leave blank if current", ""],
},
"education": {
"labels": ["Education", "Degree", "School", "University", "Institution", "Highest Education"],
"input_types": ["text", "select"],
"names": ["education", "degree", "school", "university", "institution", "edu_level"],
"placeholders": ["Bachelor's in Computer Science", "Your degree", "School name", ""],
},
"skills": {
"labels": ["Skills", "Technical Skills", "Key Skills", "Relevant Skills", "Core Competencies"],
"input_types": ["text", "textarea"],
"names": ["skills", "technical_skills", "key_skills", "competencies", "expertise"],
"placeholders": ["Python, JavaScript, AWS...", "List your skills", "Enter skills separated by commas", ""],
},
"cover_letter": {
"labels": ["Cover Letter", "Why do you want this job?", "Tell us about yourself", "Additional Information", "Message"],
"input_types": ["textarea"],
"names": ["cover_letter", "coverLetter", "message", "about", "motivation", "why_us"],
"placeholders": ["Write your cover letter here...", "Tell us why you're interested", "Additional notes", ""],
},
}
INPUT_TYPES = ["text", "email", "tel", "url", "textarea", "select", "radio", "checkbox", "file", "number", "date", "password"]What does this code define?
FIELD_CATEGORIES is the list of 18 labels your model will learn to predict. Each category represents a type of form field that appears on job applications.
FIELD_TEMPLATES maps each category to realistic variations. Real-world forms use wildly different labels for the same field ("Phone", "Cell Phone", "Contact Number"). The templates capture this diversity so the model learns to generalize.
INPUT_TYPES lists all possible HTML input types. This becomes a one-hot encoded feature later.
You won't see a change yet. The next chunk adds the functions that use these definitions to generate samples.
Seeing a syntax error?
Make sure every dictionary entry ends with a closing }, and that the outer FIELD_TEMPLATES dictionary closes with } on its own line. Nested quotes inside strings (like "Bachelor's in Computer Science") use single quotes inside double quotes, which is valid Python.
Help me fix a syntax error in my FIELD_TEMPLATES dictionary
Generate and save the training samples
With the templates defined, you need functions that randomly sample from them to produce labeled training data. The generator also injects noise (random capitalization, asterisks) so the model learns to handle messy real-world input.
- Add the following generation functions below the INPUT_TYPES list in dataset.py:
def generate_sample(category: str) -> dict:
"""Generate a single training sample for a field category."""
template = FIELD_TEMPLATES[category]
label = random.choice(template["labels"])
input_type = random.choice(template["input_types"])
name = random.choice(template["names"])
placeholder = random.choice(template["placeholders"])
# Add noise: sometimes labels have extra whitespace or punctuation
if random.random() < 0.2:
label = label + " *"
if random.random() < 0.1:
label = label.lower()
if random.random() < 0.1:
label = label.upper()
return {
"label_text": label,
"input_type": input_type,
"name_attr": name,
"placeholder": placeholder,
"target": category,
}
def generate_dataset(num_samples: int = 2000, seed: int = 42) -> list[dict]:
"""Generate a balanced training dataset."""
random.seed(seed)
samples = []
samples_per_category = num_samples // len(FIELD_CATEGORIES)
for category in FIELD_CATEGORIES:
for _ in range(samples_per_category):
samples.append(generate_sample(category))
# Fill remaining with random categories
remaining = num_samples - len(samples)
for _ in range(remaining):
category = random.choice(FIELD_CATEGORIES)
samples.append(generate_sample(category))
random.shuffle(samples)
return samples
def save_dataset(output_dir: str = "data", num_samples: int = 2000) -> tuple[Path, Path]:
"""Generate and save train/test split."""
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
samples = generate_dataset(num_samples)
split_idx = int(len(samples) * 0.8)
train_data = samples[:split_idx]
test_data = samples[split_idx:]
train_path = output_path / "form_fields_train.json"
test_path = output_path / "form_fields_test.json"
with open(train_path, "w") as f:
json.dump(train_data, f, indent=2)
with open(test_path, "w") as f:
json.dump(test_data, f, indent=2)
print(f"Training samples: {len(train_data)}")
print(f"Test samples: {len(test_data)}")
print(f"Categories: {len(FIELD_CATEGORIES)}")
print(f"Saved to: {train_path}, {test_path}")
return train_path, test_path
if __name__ == "__main__":
save_dataset()What does this code do?
generate_sample() picks random values from a category's template and injects noise. About 20% of labels get an asterisk appended (simulating required-field markers), and 10% are randomly lowercased or uppercased. This teaches the model to handle inconsistent formatting.
generate_dataset() creates a balanced dataset with equal samples per category. A fixed seed (42) ensures reproducibility so your results match what this project expects.
save_dataset() splits the data 80/20 into training and test sets, then writes them as JSON files. The test set stays untouched during training to give you an honest accuracy measurement later.
- Save dataset.py.
- Generate the dataset by running this command in your terminal:
python src/formfield_classifier/dataset.pyYou should see output confirming the split:
Training samples: 1600
Test samples: 400
Categories: 18
Saved to: data/form_fields_train.json, data/form_fields_test.json- Confirm the files exist by running:
ls data/You should see form_fields_train.json and form_fields_test.json in the output.
Seeing ModuleNotFoundError or no output?
Make sure you're running the command from the formfield-classifier/ directory (the project root), not from inside src/. The script uses relative paths to create the data/ folder.
If you see no output at all, check that your file ends with the if __name__ == "__main__": save_dataset() block.
Help me debug why my dataset generation script isn't producing output
Build the vocabulary and feature encoding
Neural networks only understand numbers. Your text fields ("First Name", "email", "john@example.com") need to be converted into fixed-length numerical vectors before the model can process them. A vocabulary maps every unique word token to an integer index.
- In dataset.py, find the line that says if __name__ == "__main__": at the bottom of the file.
- Copy the following Vocabulary class and paste it above that if __name__ line:
class Vocabulary:
"""Simple word-to-index vocabulary for text encoding."""
PAD_TOKEN = "<PAD>"
UNK_TOKEN = "<UNK>"
def __init__(self):
self.word2idx: dict[str, int] = {self.PAD_TOKEN: 0, self.UNK_TOKEN: 1}
self.idx2word: dict[int, str] = {0: self.PAD_TOKEN, 1: self.UNK_TOKEN}
self.word_count: dict[str, int] = {}
def build_from_samples(self, samples: list[dict], min_freq: int = 1):
"""Build vocabulary from dataset samples."""
for sample in samples:
for field in ["label_text", "name_attr", "placeholder"]:
tokens = self.tokenize(sample[field])
for token in tokens:
self.word_count[token] = self.word_count.get(token, 0) + 1
for word, count in self.word_count.items():
if count >= min_freq and word not in self.word2idx:
idx = len(self.word2idx)
self.word2idx[word] = idx
self.idx2word[idx] = word
print(f"Vocabulary size: {len(self.word2idx)}")
def tokenize(self, text: str) -> list[str]:
"""Simple whitespace + punctuation tokenizer."""
text = text.lower().strip()
tokens = []
current = ""
for char in text:
if char.isalnum():
current += char
else:
if current:
tokens.append(current)
current = ""
if char.strip():
tokens.append(char)
if current:
tokens.append(current)
return tokens
def encode(self, text: str, max_length: int = 10) -> list[int]:
"""Encode text to list of token indices."""
tokens = self.tokenize(text)[:max_length]
indices = [self.word2idx.get(t, self.word2idx[self.UNK_TOKEN]) for t in tokens]
# Pad to max_length
indices += [self.word2idx[self.PAD_TOKEN]] * (max_length - len(indices))
return indices
def save(self, path: str):
"""Save vocabulary to JSON."""
with open(path, "w") as f:
json.dump({"word2idx": self.word2idx}, f)
def load(self, path: str):
"""Load vocabulary from JSON."""
with open(path, "r") as f:
data = json.load(f)
self.word2idx = data["word2idx"]
self.idx2word = {int(v): k for k, v in self.word2idx.items()}What does the Vocabulary class do?
tokenize() splits text into lowercase word tokens. "First Name *" becomes ["first", "name", "*"]. This normalization helps the model treat "First Name" and "first name" as the same tokens.
build_from_samples() scans every text field in the training data and assigns a unique integer to each word. Special tokens <PAD> (index 0) and <UNK> (index 1) handle padding and unseen words.
encode() converts a text string into a fixed-length list of integers. Short sequences get padded with zeros. This fixed shape is what PyTorch tensors require.
- Save dataset.py.
- Verify the vocabulary builds correctly and produces valid tensors by running:
python -c "
import json, torch, sys
sys.path.insert(0, 'src')
from formfield_classifier.dataset import Vocabulary
with open('data/form_fields_train.json') as f:
train_data = json.load(f)
vocab = Vocabulary()
vocab.build_from_samples(train_data)
# Encode a sample and convert to a PyTorch tensor
encoded = vocab.encode('First Name')
tensor = torch.tensor([encoded], dtype=torch.long)
print(f'Encoded: {encoded}')
print(f'Tensor shape: {tensor.shape}')
print(f'Device available: {torch.cuda.is_available()}')
"You should see the vocabulary size printed (around 200-300 tokens), followed by the encoded integer list, a tensor shape of torch.Size([1, 10]), and Device available: True.
Seeing an import error for torch?
Make sure you installed PyTorch in Step 1. Run pip list | grep torch to confirm it's installed. If the sys.path.insert line gives you trouble, confirm you're running the command from the formfield-classifier/ project root.
Help me debug an import error when verifying my Vocabulary class
✔️ Awesome, I've got everything!
Great. Double check you've saved dataset.py and that both JSON files exist in your data/ folder.
ⓧ I'd like to double check the full code
Here's the complete src/formfield_classifier/dataset.py file:
"""Training dataset generation and loading for form field classification."""
import json
import random
from pathlib import Path
FIELD_CATEGORIES = [
"first_name", "last_name", "email", "phone", "linkedin",
"portfolio_url", "address", "city", "state", "zip_code",
"country", "company", "job_title", "start_date", "end_date",
"education", "skills", "cover_letter",
]
FIELD_TEMPLATES = {
"first_name": {
"labels": ["First Name", "First name", "Given Name", "Your first name", "Name (First)", "Forename"],
"input_types": ["text"],
"names": ["first_name", "firstName", "fname", "given_name", "first-name", "givenName"],
"placeholders": ["John", "Enter first name", "First name", "Your first name", ""],
},
"last_name": {
"labels": ["Last Name", "Last name", "Surname", "Family Name", "Name (Last)", "Your last name"],
"input_types": ["text"],
"names": ["last_name", "lastName", "lname", "surname", "family_name", "familyName"],
"placeholders": ["Doe", "Enter last name", "Last name", "Your surname", ""],
},
"email": {
"labels": ["Email", "Email Address", "E-mail", "Your Email", "Work Email", "Personal Email"],
"input_types": ["email", "text"],
"names": ["email", "emailAddress", "user_email", "applicant_email", "e-mail", "mail"],
"placeholders": ["john@example.com", "Enter email", "you@company.com", "Email address", ""],
},
"phone": {
"labels": ["Phone", "Phone Number", "Mobile", "Cell Phone", "Contact Number", "Telephone"],
"input_types": ["tel", "text"],
"names": ["phone", "phoneNumber", "mobile", "telephone", "cell", "phone_number"],
"placeholders": ["+1 (555) 123-4567", "Enter phone number", "Phone number", "Your phone", ""],
},
"linkedin": {
"labels": ["LinkedIn", "LinkedIn URL", "LinkedIn Profile", "Your LinkedIn", "LinkedIn Profile URL"],
"input_types": ["url", "text"],
"names": ["linkedin", "linkedinUrl", "linkedin_url", "linkedin_profile", "social_linkedin"],
"placeholders": ["https://linkedin.com/in/yourname", "LinkedIn URL", "Your LinkedIn profile", ""],
},
"portfolio_url": {
"labels": ["Portfolio", "Website", "Personal Website", "Portfolio URL", "Your Website", "GitHub"],
"input_types": ["url", "text"],
"names": ["portfolio", "website", "personal_url", "portfolio_url", "github", "url"],
"placeholders": ["https://yoursite.com", "Your website URL", "Portfolio link", ""],
},
"address": {
"labels": ["Address", "Street Address", "Home Address", "Address Line 1", "Street", "Mailing Address"],
"input_types": ["text"],
"names": ["address", "street_address", "address1", "street", "home_address", "addr"],
"placeholders": ["123 Main St", "Enter your address", "Street address", ""],
},
"city": {
"labels": ["City", "City/Town", "Your City", "Town", "Municipality"],
"input_types": ["text"],
"names": ["city", "town", "municipality", "city_name", "applicant_city"],
"placeholders": ["New York", "Enter city", "City", "Your city", ""],
},
"state": {
"labels": ["State", "State/Province", "Province", "Region", "State/Region"],
"input_types": ["text", "select"],
"names": ["state", "province", "region", "state_province", "applicant_state"],
"placeholders": ["California", "Select state", "State", "Province/State", ""],
},
"zip_code": {
"labels": ["Zip Code", "ZIP", "Postal Code", "Zip/Postal Code", "Postcode"],
"input_types": ["text"],
"names": ["zip", "zip_code", "postal_code", "zipCode", "postcode", "postalCode"],
"placeholders": ["10001", "Enter zip code", "Zip code", "Postal code", ""],
},
"country": {
"labels": ["Country", "Country/Region", "Your Country", "Nation", "Country of Residence"],
"input_types": ["select", "text"],
"names": ["country", "nation", "country_code", "residence_country", "applicant_country"],
"placeholders": ["United States", "Select country", "Country", ""],
},
"company": {
"labels": ["Company", "Company Name", "Employer", "Organization", "Current Company", "Previous Employer"],
"input_types": ["text"],
"names": ["company", "employer", "organization", "company_name", "org", "firm"],
"placeholders": ["Acme Corp", "Company name", "Your current employer", "Organization name", ""],
},
"job_title": {
"labels": ["Job Title", "Title", "Position", "Role", "Current Title", "Position Title"],
"input_types": ["text"],
"names": ["job_title", "title", "position", "role", "current_title", "jobTitle"],
"placeholders": ["Software Engineer", "Your job title", "Position title", "Current role", ""],
},
"start_date": {
"labels": ["Start Date", "From", "Date Started", "Begin Date", "Employment Start"],
"input_types": ["date", "text"],
"names": ["start_date", "startDate", "from_date", "date_started", "begin_date"],
"placeholders": ["MM/YYYY", "Start date", "When did you start?", "2020-01", ""],
},
"end_date": {
"labels": ["End Date", "To", "Date Ended", "End", "Employment End", "Currently Working"],
"input_types": ["date", "text"],
"names": ["end_date", "endDate", "to_date", "date_ended", "end", "until"],
"placeholders": ["MM/YYYY", "End date", "Present", "Leave blank if current", ""],
},
"education": {
"labels": ["Education", "Degree", "School", "University", "Institution", "Highest Education"],
"input_types": ["text", "select"],
"names": ["education", "degree", "school", "university", "institution", "edu_level"],
"placeholders": ["Bachelor's in Computer Science", "Your degree", "School name", ""],
},
"skills": {
"labels": ["Skills", "Technical Skills", "Key Skills", "Relevant Skills", "Core Competencies"],
"input_types": ["text", "textarea"],
"names": ["skills", "technical_skills", "key_skills", "competencies", "expertise"],
"placeholders": ["Python, JavaScript, AWS...", "List your skills", "Enter skills separated by commas", ""],
},
"cover_letter": {
"labels": ["Cover Letter", "Why do you want this job?", "Tell us about yourself", "Additional Information", "Message"],
"input_types": ["textarea"],
"names": ["cover_letter", "coverLetter", "message", "about", "motivation", "why_us"],
"placeholders": ["Write your cover letter here...", "Tell us why you're interested", "Additional notes", ""],
},
}
INPUT_TYPES = ["text", "email", "tel", "url", "textarea", "select", "radio", "checkbox", "file", "number", "date", "password"]
def generate_sample(category: str) -> dict:
"""Generate a single training sample for a field category."""
template = FIELD_TEMPLATES[category]
label = random.choice(template["labels"])
input_type = random.choice(template["input_types"])
name = random.choice(template["names"])
placeholder = random.choice(template["placeholders"])
if random.random() < 0.2:
label = label + " *"
if random.random() < 0.1:
label = label.lower()
if random.random() < 0.1:
label = label.upper()
return {
"label_text": label,
"input_type": input_type,
"name_attr": name,
"placeholder": placeholder,
"target": category,
}
def generate_dataset(num_samples: int = 2000, seed: int = 42) -> list[dict]:
"""Generate a balanced training dataset."""
random.seed(seed)
samples = []
samples_per_category = num_samples // len(FIELD_CATEGORIES)
for category in FIELD_CATEGORIES:
for _ in range(samples_per_category):
samples.append(generate_sample(category))
remaining = num_samples - len(samples)
for _ in range(remaining):
category = random.choice(FIELD_CATEGORIES)
samples.append(generate_sample(category))
random.shuffle(samples)
return samples
def save_dataset(output_dir: str = "data", num_samples: int = 2000) -> tuple[Path, Path]:
"""Generate and save train/test split."""
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
samples = generate_dataset(num_samples)
split_idx = int(len(samples) * 0.8)
train_data = samples[:split_idx]
test_data = samples[split_idx:]
train_path = output_path / "form_fields_train.json"
test_path = output_path / "form_fields_test.json"
with open(train_path, "w") as f:
json.dump(train_data, f, indent=2)
with open(test_path, "w") as f:
json.dump(test_data, f, indent=2)
print(f"Training samples: {len(train_data)}")
print(f"Test samples: {len(test_data)}")
print(f"Categories: {len(FIELD_CATEGORIES)}")
print(f"Saved to: {train_path}, {test_path}")
return train_path, test_path
class Vocabulary:
"""Simple word-to-index vocabulary for text encoding."""
PAD_TOKEN = "<PAD>"
UNK_TOKEN = "<UNK>"
def __init__(self):
self.word2idx: dict[str, int] = {self.PAD_TOKEN: 0, self.UNK_TOKEN: 1}
self.idx2word: dict[int, str] = {0: self.PAD_TOKEN, 1: self.UNK_TOKEN}
self.word_count: dict[str, int] = {}
def build_from_samples(self, samples: list[dict], min_freq: int = 1):
"""Build vocabulary from dataset samples."""
for sample in samples:
for field in ["label_text", "name_attr", "placeholder"]:
tokens = self.tokenize(sample[field])
for token in tokens:
self.word_count[token] = self.word_count.get(token, 0) + 1
for word, count in self.word_count.items():
if count >= min_freq and word not in self.word2idx:
idx = len(self.word2idx)
self.word2idx[word] = idx
self.idx2word[idx] = word
print(f"Vocabulary size: {len(self.word2idx)}")
def tokenize(self, text: str) -> list[str]:
"""Simple whitespace + punctuation tokenizer."""
text = text.lower().strip()
tokens = []
current = ""
for char in text:
if char.isalnum():
current += char
else:
if current:
tokens.append(current)
current = ""
if char.strip():
tokens.append(char)
if current:
tokens.append(current)
return tokens
def encode(self, text: str, max_length: int = 10) -> list[int]:
"""Encode text to list of token indices."""
tokens = self.tokenize(text)[:max_length]
indices = [self.word2idx.get(t, self.word2idx[self.UNK_TOKEN]) for t in tokens]
indices += [self.word2idx[self.PAD_TOKEN]] * (max_length - len(indices))
return indices
def save(self, path: str):
"""Save vocabulary to JSON."""
with open(path, "w") as f:
json.dump({"word2idx": self.word2idx}, f)
def load(self, path: str):
"""Load vocabulary from JSON."""
with open(path, "r") as f:
data = json.load(f)
self.word2idx = data["word2idx"]
self.idx2word = {int(v): k for k, v in self.word2idx.items()}
if __name__ == "__main__":
save_dataset()Why synthetic data instead of real forms?
In production, you'd scrape real job application forms from platforms like Greenhouse, Lever, and Workday. Synthetic data teaches the full ML pipeline without requiring web scraping infrastructure.
The template approach still produces a model that generalizes well because it captures the real variation patterns: multiple label phrasings, different input types, noisy formatting. When you later feed it real form HTML, the model has already seen these patterns.
Your training data pipeline is complete. Next up, you'll design the neural network architecture that learns to classify these form fields from the numerical features you just created.
Build the Neural Network Architecture
You now have a labeled dataset of 2000+ form field samples split into training and test sets. The next challenge is building a model that can learn from this data.
In this step, you will design a PyTorch neural network from scratch that takes raw text metadata (labels, names, placeholders) and predicts which of the 18 field categories a form field belongs to.
In this step, get ready to:
- Define the FormFieldClassifier model class with embedding and MLP layers.
- Implement the forward pass with bag-of-embeddings encoding and feature concatenation.
- Verify the model produces the correct output shape and parameter count.
Define the FormFieldClassifier class
The classifier needs to handle multiple text inputs (label, name attribute, placeholder) plus a categorical input (HTML input type). Your architecture will use a shared embedding layer to convert text tokens into dense vectors, then concatenate all features before passing them through fully connected layers.
- Create a new file called model.py inside src/formfield_classifier/ and add the imports and class definition:
"""PyTorch neural network architecture for form field classification."""
import torch
import torch.nn as nn
from .dataset import INPUT_TYPES, FIELD_CATEGORIES
class FormFieldClassifier(nn.Module):
"""Multi-layer neural network for classifying form field types.
Architecture:
- Shared embedding layer for text tokens
- Bag-of-embeddings encoder (average pooling over token embeddings)
- Three text features (label, name, placeholder) encoded independently
- One-hot input type concatenated with text encodings
- MLP head with dropout for classification
"""
def __init__(
self,
vocab_size: int,
embed_dim: int = 32,
hidden_dim: int = 128,
num_classes: int = len(FIELD_CATEGORIES),
num_input_types: int = len(INPUT_TYPES),
dropout: float = 0.3,
):
super().__init__()
self.num_input_types = num_input_types
self.embed_dim = embed_dim
# Shared embedding for all text fields
self.embedding = nn.Embedding(
num_embeddings=vocab_size,
embedding_dim=embed_dim,
padding_idx=0,
)
# Input feature size: 3 text fields (each embed_dim) + one-hot input type
input_size = (3 * embed_dim) + num_input_types
# MLP classification head
self.classifier = nn.Sequential(
nn.Linear(input_size, hidden_dim),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, hidden_dim // 2),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim // 2, num_classes),
)What does this architecture do?
The FormFieldClassifier subclasses nn.Module, which is the base class for all PyTorch neural networks. The constructor defines three key components:
- self.embedding is a lookup table that converts word indices into 32-dimensional dense vectors. All three text fields share this single embedding layer.
- input_size calculates the total feature vector width: three text fields (each 32 dims) plus the one-hot input type vector (12 dims) = 108.
- self.classifier is a three-layer MLP (108 to 128 to 64 to 18) with ReLU activations and 30% dropout for regularization.
Why bag-of-embeddings instead of pretrained word vectors?
Pretrained embeddings like GloVe or Word2Vec are trained on general text. Form field metadata uses a specialized vocabulary (attribute names, HTML types, abbreviated labels) that general embeddings handle poorly.
By learning embeddings from scratch, the model discovers that tokens like "fname" and "first" are closely related in this domain. The model is also tiny (under 100KB), which makes it ideal for fast inference in a browser extension backend.
Implement the forward pass
The forward pass defines how data flows through the network. You need two methods: one that encodes a batch of text tokens into fixed-size vectors, and one that orchestrates the full prediction pipeline.
- Add the encode_text method inside the FormFieldClassifier class, below the __init__ method:
def encode_text(self, token_indices: torch.Tensor) -> torch.Tensor:
"""Encode text tokens to a fixed-size vector via bag-of-embeddings.
Args:
token_indices: (batch_size, max_length) tensor of token indices
Returns:
(batch_size, embed_dim) averaged embeddings
"""
embedded = self.embedding(token_indices) # (batch, max_len, embed_dim)
# Create mask for non-padding tokens
mask = (token_indices != 0).unsqueeze(-1).float() # (batch, max_len, 1)
# Average pooling over non-padding tokens
summed = (embedded * mask).sum(dim=1) # (batch, embed_dim)
lengths = mask.sum(dim=1).clamp(min=1) # (batch, 1) avoid div by zero
averaged = summed / lengths # (batch, embed_dim)
return averagedWhat does encode_text do?
This method converts a batch of variable-length token sequences into fixed-size vectors using masked average pooling:
- The mask zeros out padding tokens (index 0) so they don't contribute to the average.
- lengths.clamp(min=1) prevents division by zero when a text field is completely empty.
- The result is a single vector per text field that summarizes all its tokens, regardless of how many tokens the field had.
- Add the forward method below encode_text:
def forward(
self,
label_tokens: torch.Tensor,
name_tokens: torch.Tensor,
placeholder_tokens: torch.Tensor,
input_type_onehot: torch.Tensor,
) -> torch.Tensor:
"""Forward pass through the network.
Args:
label_tokens: (batch, max_len) encoded label text
name_tokens: (batch, max_len) encoded name attribute
placeholder_tokens: (batch, max_len) encoded placeholder text
input_type_onehot: (batch, num_input_types) one-hot input type
Returns:
(batch, num_classes) logits for each field category
"""
label_encoded = self.encode_text(label_tokens)
name_encoded = self.encode_text(name_tokens)
placeholder_encoded = self.encode_text(placeholder_tokens)
# Concatenate all features
features = torch.cat(
[label_encoded, name_encoded, placeholder_encoded, input_type_onehot],
dim=1,
)
logits = self.classifier(features)
return logitsWhat does the forward method do?
The forward method is what PyTorch calls automatically when you pass data through the model. It encodes each text field independently through the shared embedding layer, concatenates all three encoded vectors with the one-hot input type, and passes the combined feature vector through the MLP classifier.
The output is raw logits (unnormalized scores) for each of the 18 categories. During training, these logits feed into the loss function. During inference, you apply softmax to get probabilities.
Add the factory function and verify the model
- Below the class (outside it, at the module level), add the create_model factory function:
def create_model(vocab_size: int, device: str = "cuda") -> FormFieldClassifier:
"""Factory function to create and move model to device."""
model = FormFieldClassifier(vocab_size=vocab_size)
model = model.to(device)
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model created on {device}")
print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")
return modelWhat does create_model do?
This factory function handles model instantiation and device placement in one call. It also prints the parameter count so you can confirm the model is the expected size. Moving the model to the CUDA device ensures all computations run on your GPU.
- Save model.py.
- Verify the model by running this test script from your formfield-classifier directory:
python -c "
import torch
from src.formfield_classifier.model import create_model
model = create_model(vocab_size=200, device='cuda')
# Create a dummy batch of 4 samples
batch_size = 4
max_len = 10
label = torch.randint(0, 200, (batch_size, max_len)).cuda()
name = torch.randint(0, 200, (batch_size, max_len)).cuda()
placeholder = torch.randint(0, 200, (batch_size, max_len)).cuda()
input_type = torch.zeros(batch_size, 12).cuda()
input_type[:, 0] = 1.0
output = model(label, name, placeholder, input_type)
print(f'Output shape: {output.shape}')
"You should see output similar to:
Model created on cuda
Total parameters: 25,938
Trainable parameters: 25,938
Output shape: torch.Size([4, 18])What should I see?
The exact parameter count depends on your vocabulary size. With a vocab of 200, you get around 25K parameters. With the full vocabulary from your dataset (which will be larger), expect 50K-100K parameters. The important confirmation is that the output shape is (batch_size, 18), meaning the model produces one score per category for each sample in the batch.
Seeing a CUDA error?
- If you see a RuntimeError about CUDA not being available, change device='cuda' to device='cpu' in the test command. Then revisit Step 1 to confirm your GPU setup.
- If you see an ImportError, make sure you are running the command from the formfield-classifier root directory (not from inside src/).
- If the output shape is wrong, double check that FIELD_CATEGORIES in dataset.py has exactly 18 items.
Why isn't my FormFieldClassifier model producing the correct output shape?
✔️ Awesome, I've got everything!
Double check you have saved src/formfield_classifier/model.py.
ⓧ I'd like to double check the full code
"""PyTorch neural network architecture for form field classification."""
import torch
import torch.nn as nn
from .dataset import INPUT_TYPES, FIELD_CATEGORIES
class FormFieldClassifier(nn.Module):
def __init__(
self,
vocab_size: int,
embed_dim: int = 32,
hidden_dim: int = 128,
num_classes: int = len(FIELD_CATEGORIES),
num_input_types: int = len(INPUT_TYPES),
dropout: float = 0.3,
):
super().__init__()
self.num_input_types = num_input_types
self.embed_dim = embed_dim
self.embedding = nn.Embedding(
num_embeddings=vocab_size,
embedding_dim=embed_dim,
padding_idx=0,
)
input_size = (3 * embed_dim) + num_input_types
self.classifier = nn.Sequential(
nn.Linear(input_size, hidden_dim),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, hidden_dim // 2),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim // 2, num_classes),
)
def encode_text(self, token_indices: torch.Tensor) -> torch.Tensor:
embedded = self.embedding(token_indices)
mask = (token_indices != 0).unsqueeze(-1).float()
summed = (embedded * mask).sum(dim=1)
lengths = mask.sum(dim=1).clamp(min=1)
averaged = summed / lengths
return averaged
def forward(
self,
label_tokens: torch.Tensor,
name_tokens: torch.Tensor,
placeholder_tokens: torch.Tensor,
input_type_onehot: torch.Tensor,
) -> torch.Tensor:
label_encoded = self.encode_text(label_tokens)
name_encoded = self.encode_text(name_tokens)
placeholder_encoded = self.encode_text(placeholder_tokens)
features = torch.cat(
[label_encoded, name_encoded, placeholder_encoded, input_type_onehot],
dim=1,
)
logits = self.classifier(features)
return logits
def create_model(vocab_size: int, device: str = "cuda") -> FormFieldClassifier:
model = FormFieldClassifier(vocab_size=vocab_size)
model = model.to(device)
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model created on {device}")
print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")
return modelYour neural network architecture is locked in. Next up, you will write the training loop that feeds your dataset through this model, computes loss, and updates the weights using GPU acceleration.
Train the Model
You have a neural network architecture that accepts form field metadata and outputs classification logits. The architecture is sound, but it has never seen real data.
In this step, you will write the full training pipeline: a PyTorch Dataset class that encodes raw samples into tensors, a training loop that runs on your GPU, evaluation functions that measure accuracy, and a registration step that stores the trained model in DataEngineX's model registry.
In this step, get ready to:
- Implement a PyTorch Dataset and GPU-accelerated training loop.
- Add evaluation functions that report per-class accuracy.
- Register the trained model in DataEngineX's ModelRegistry.
Implement the training loop
The training loop needs data in tensor format. A custom Dataset class handles the conversion from raw JSON samples to padded token indices, one-hot vectors, and integer targets that PyTorch can batch and send to your GPU.
- Create a new file src/formfield_classifier/train.py and add the imports and FormFieldDataset class by pasting the following code:
"""Training loop with GPU acceleration and DataEngineX integration."""
import json
import time
from datetime import datetime, timezone
from pathlib import Path
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from dataenginex.engine import DexEngine
from dataenginex.ml import ModelRegistry, ModelArtifact, ModelStage
from .dataset import FIELD_CATEGORIES, INPUT_TYPES, Vocabulary
from .model import FormFieldClassifier, create_model
class FormFieldDataset(Dataset):
"""PyTorch Dataset for form field samples."""
def __init__(self, samples: list[dict], vocab: Vocabulary, max_length: int = 10):
self.samples = samples
self.vocab = vocab
self.max_length = max_length
self.category_to_idx = {cat: i for i, cat in enumerate(FIELD_CATEGORIES)}
self.input_type_to_idx = {t: i for i, t in enumerate(INPUT_TYPES)}
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, idx: int) -> dict:
sample = self.samples[idx]
# Encode each text field into fixed-length token indices
label_tokens = torch.tensor(
self.vocab.encode(sample["label_text"], self.max_length), dtype=torch.long
)
name_tokens = torch.tensor(
self.vocab.encode(sample["name_attr"], self.max_length), dtype=torch.long
)
placeholder_tokens = torch.tensor(
self.vocab.encode(sample["placeholder"], self.max_length), dtype=torch.long
)
# One-hot encode the HTML input type
input_type_vec = torch.zeros(len(INPUT_TYPES), dtype=torch.float)
type_idx = self.input_type_to_idx.get(sample["input_type"], 0)
input_type_vec[type_idx] = 1.0
# Integer target for CrossEntropyLoss
target = torch.tensor(
self.category_to_idx[sample["target"]], dtype=torch.long
)
return {
"label_tokens": label_tokens,
"name_tokens": name_tokens,
"placeholder_tokens": placeholder_tokens,
"input_type": input_type_vec,
"target": target,
}What does this code do?
The imports pull in PyTorch's training utilities, your model and dataset modules, and the DataEngineX classes you will use later for model registration.
FormFieldDataset converts each raw JSON sample into tensors. It encodes the three text fields (label_text, name_attr, placeholder) into fixed-length integer sequences using the vocabulary. The HTML input type becomes a one-hot vector. The target category becomes a single integer index for CrossEntropyLoss.
- You won't see output yet. The train_model() function uses this class next.
Now add the train_model() function below the FormFieldDataset class. This is the core training loop that loads data, builds the vocabulary, creates DataLoaders, and iterates over 50 epochs on your GPU.
- Add the following function to train.py below the class you just wrote:
def train_model(
data_dir: str = "data",
model_dir: str = "models",
epochs: int = 50,
batch_size: int = 64,
learning_rate: float = 0.001,
device: str = "cuda",
) -> dict:
"""Train the form field classifier.
Returns:
Dictionary with training results and metrics.
"""
data_path = Path(data_dir)
model_path = Path(model_dir)
model_path.mkdir(parents=True, exist_ok=True)
# Load training and test data from JSON
with open(data_path / "form_fields_train.json") as f:
train_data = json.load(f)
with open(data_path / "form_fields_test.json") as f:
test_data = json.load(f)
# Build vocabulary from training samples and save it
vocab = Vocabulary()
vocab.build_from_samples(train_data)
vocab.save(str(model_path / "vocabulary.json"))
# Create datasets and data loaders
train_dataset = FormFieldDataset(train_data, vocab)
test_dataset = FormFieldDataset(test_data, vocab)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# Create model on GPU
model = create_model(vocab_size=len(vocab.word2idx), device=device)
# Loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# Training loop
print(f"\nTraining for {epochs} epochs on {device}...")
print("-" * 60)
best_accuracy = 0.0
train_start = time.time()
for epoch in range(epochs):
model.train()
total_loss = 0.0
correct = 0
total = 0
for batch in train_loader:
# Move all tensors to GPU
label_tokens = batch["label_tokens"].to(device)
name_tokens = batch["name_tokens"].to(device)
placeholder_tokens = batch["placeholder_tokens"].to(device)
input_type = batch["input_type"].to(device)
targets = batch["target"].to(device)
# Forward pass
optimizer.zero_grad()
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type)
loss = criterion(logits, targets)
# Backward pass and weight update
loss.backward()
optimizer.step()
# Track metrics
total_loss += loss.item()
predictions = logits.argmax(dim=1)
correct += (predictions == targets).sum().item()
total += targets.size(0)
train_accuracy = correct / total
avg_loss = total_loss / len(train_loader)
# Evaluate on test set every 10 epochs
if (epoch + 1) % 10 == 0 or epoch == epochs - 1:
test_accuracy = evaluate(model, test_loader, device)
print(
f"Epoch {epoch + 1:3d}/{epochs} | "
f"Loss: {avg_loss:.4f} | "
f"Train Acc: {train_accuracy:.4f} | "
f"Test Acc: {test_accuracy:.4f}"
)
# Save best model checkpoint
if test_accuracy > best_accuracy:
best_accuracy = test_accuracy
torch.save(model.state_dict(), model_path / "form_field_classifier_v1.pt")
train_duration = time.time() - train_start
print("-" * 60)
print(f"Training complete in {train_duration:.1f}s")
print(f"Best test accuracy: {best_accuracy:.4f}")
# Final detailed per-class evaluation
print("\nPer-class accuracy:")
class_report = evaluate_detailed(model, test_loader, device)
for category, acc in class_report.items():
print(f" {category:20s}: {acc:.4f}")
results = {
"best_accuracy": best_accuracy,
"final_loss": avg_loss,
"epochs": epochs,
"training_time_seconds": train_duration,
"vocab_size": len(vocab.word2idx),
"model_path": str(model_path / "form_field_classifier_v1.pt"),
"class_report": class_report,
}
return resultsWhat does the training loop do?
The function loads your JSON datasets, builds a vocabulary, and wraps everything in DataLoaders that batch 64 samples at a time. It creates the model on your CUDA device and uses CrossEntropyLoss (the standard loss for multi-class classification) with the Adam optimizer.
Each epoch iterates over all batches: move tensors to GPU, compute predictions, calculate loss, backpropagate gradients, and update weights. Every 10 epochs, it evaluates on the held-out test set and saves the model if accuracy improves.
- You won't see output yet. The evaluate() functions referenced in the loop come next.
Implement evaluation
The training loop calls two evaluation helpers. One computes overall accuracy across the test set. The other breaks accuracy down per field category so you can see which categories the model handles well.
- Add both evaluation functions below train_model() in train.py:
def evaluate(model: FormFieldClassifier, loader: DataLoader, device: str) -> float:
"""Evaluate model accuracy on a data loader."""
model.eval()
correct = 0
total = 0
with torch.no_grad():
for batch in loader:
label_tokens = batch["label_tokens"].to(device)
name_tokens = batch["name_tokens"].to(device)
placeholder_tokens = batch["placeholder_tokens"].to(device)
input_type = batch["input_type"].to(device)
targets = batch["target"].to(device)
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type)
predictions = logits.argmax(dim=1)
correct += (predictions == targets).sum().item()
total += targets.size(0)
return correct / total
def evaluate_detailed(
model: FormFieldClassifier, loader: DataLoader, device: str
) -> dict[str, float]:
"""Per-class accuracy evaluation."""
model.eval()
class_correct: dict[int, int] = {i: 0 for i in range(len(FIELD_CATEGORIES))}
class_total: dict[int, int] = {i: 0 for i in range(len(FIELD_CATEGORIES))}
with torch.no_grad():
for batch in loader:
label_tokens = batch["label_tokens"].to(device)
name_tokens = batch["name_tokens"].to(device)
placeholder_tokens = batch["placeholder_tokens"].to(device)
input_type = batch["input_type"].to(device)
targets = batch["target"].to(device)
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type)
predictions = logits.argmax(dim=1)
for pred, target in zip(predictions, targets):
t = target.item()
class_total[t] += 1
if pred.item() == t:
class_correct[t] += 1
report = {}
for i, category in enumerate(FIELD_CATEGORIES):
if class_total[i] > 0:
report[category] = class_correct[i] / class_total[i]
else:
report[category] = 0.0
return reportWhy two evaluation functions?
evaluate() gives you a single overall accuracy number to track during training. evaluate_detailed() breaks accuracy down by category so you can spot which field types the model struggles with (often ambiguous ones like cover_letter vs skills).
Both use torch.no_grad() to disable gradient computation during evaluation. This saves memory and speeds up inference since you only need forward passes here.
- You won't see output yet. The next section adds model registration and runs everything.
Integrate with DataEngineX and run training
The final piece connects your trained model to the DataEngineX ecosystem. The register_model() function stores the model artifact with metadata (accuracy, training time, vocab size) and promotes it to production if it meets your accuracy threshold.
- Add the registration function and the __main__ block at the bottom of train.py:
def register_model(results: dict, dex_config: str = "dex.yaml"):
"""Register trained model in DataEngineX ModelRegistry."""
engine = DexEngine(dex_config)
registry = ModelRegistry()
# Create model artifact with training metadata
artifact = ModelArtifact(
name="form_field_classifier",
version="1.0.0",
stage=ModelStage.STAGING,
metadata={
"framework": "pytorch",
"accuracy": results["best_accuracy"],
"training_time": results["training_time_seconds"],
"vocab_size": results["vocab_size"],
"num_classes": len(FIELD_CATEGORIES),
"model_path": results["model_path"],
},
)
# Register at STAGING first
registry.register(artifact, name="form_field_classifier", stage=ModelStage.STAGING)
print(f"Model registered: form_field_classifier v1.0.0 (STAGING)")
# Promote to production if accuracy meets threshold
if results["best_accuracy"] >= 0.85:
registry.register(artifact, name="form_field_classifier", stage=ModelStage.PRODUCTION)
print("Model promoted to PRODUCTION (accuracy >= 85%)")
return artifact
if __name__ == "__main__":
results = train_model()
register_model(results)What does model registration do?
ModelRegistry stores versioned model artifacts with metadata in your local .dex/store.duckdb database. The artifact records the framework, accuracy, training time, and file path so you can retrieve and serve the model later.
The two-stage pattern (STAGING then PRODUCTION) is a common MLOps practice. You register first, then promote only if the model passes your quality gate (85% accuracy in this case).
- Save src/formfield_classifier/train.py.
- Run the training script by executing this command from your formfield-classifier folder:
python -m formfield_classifier.trainTraining runs for 50 epochs on your Quadro T2000. You should see output printing every 10 epochs with loss decreasing and accuracy climbing above 85%.
What should I see?
The script prints vocabulary size, model parameter counts, then epoch-by-epoch metrics. After about 1-2 minutes, you should see something like:
Vocabulary size: 210
Model created on cuda
Total parameters: 52,946
Trainable parameters: 52,946
Training for 50 epochs on cuda...
------------------------------------------------------------
Epoch 10/50 | Loss: 0.8234 | Train Acc: 0.7812 | Test Acc: 0.8025
Epoch 20/50 | Loss: 0.3156 | Train Acc: 0.9103 | Test Acc: 0.8750
Epoch 30/50 | Loss: 0.1542 | Train Acc: 0.9621 | Test Acc: 0.9125
...
------------------------------------------------------------
Training complete in 45.2s
Best test accuracy: 0.9250
Model registered: form_field_classifier v1.0.0 (STAGING)
Model promoted to PRODUCTION (accuracy >= 85%)Your exact numbers will differ slightly due to random initialization, but test accuracy should exceed 85%. The per-class report shows which categories perform best.
Training not starting or erroring out?
- If you see an ImportError, make sure you are running from the formfield-classifier directory (one level above src/). The -m flag requires the package to be discoverable from your current directory.
- If you see a CUDA out of memory error, close any other GPU-using processes (like a Jupyter notebook) and retry. Your model is small but other programs may be using VRAM.
- If training data is not found, confirm you generated the dataset in Step 2 and that data/form_fields_train.json exists.
Help me debug my training script
- Confirm you see the final line showing the model promoted to PRODUCTION.
- Verify the model file was saved by running:
ls models/- You should see form_field_classifier_v1.pt and vocabulary.json in the output.
✔️ Awesome, I've got everything!
Great. Double check you have saved train.py and that both models/form_field_classifier_v1.pt and models/vocabulary.json exist on disk.
ⓧ I'd like to double check the full code
Here is the complete src/formfield_classifier/train.py file for comparison:
"""Training loop with GPU acceleration and DataEngineX integration."""
import json
import time
from datetime import datetime, timezone
from pathlib import Path
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from dataenginex.engine import DexEngine
from dataenginex.ml import ModelRegistry, ModelArtifact, ModelStage
from .dataset import FIELD_CATEGORIES, INPUT_TYPES, Vocabulary
from .model import FormFieldClassifier, create_model
class FormFieldDataset(Dataset):
"""PyTorch Dataset for form field samples."""
def __init__(self, samples: list[dict], vocab: Vocabulary, max_length: int = 10):
self.samples = samples
self.vocab = vocab
self.max_length = max_length
self.category_to_idx = {cat: i for i, cat in enumerate(FIELD_CATEGORIES)}
self.input_type_to_idx = {t: i for i, t in enumerate(INPUT_TYPES)}
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, idx: int) -> dict:
sample = self.samples[idx]
label_tokens = torch.tensor(
self.vocab.encode(sample["label_text"], self.max_length), dtype=torch.long
)
name_tokens = torch.tensor(
self.vocab.encode(sample["name_attr"], self.max_length), dtype=torch.long
)
placeholder_tokens = torch.tensor(
self.vocab.encode(sample["placeholder"], self.max_length), dtype=torch.long
)
input_type_vec = torch.zeros(len(INPUT_TYPES), dtype=torch.float)
type_idx = self.input_type_to_idx.get(sample["input_type"], 0)
input_type_vec[type_idx] = 1.0
target = torch.tensor(
self.category_to_idx[sample["target"]], dtype=torch.long
)
return {
"label_tokens": label_tokens,
"name_tokens": name_tokens,
"placeholder_tokens": placeholder_tokens,
"input_type": input_type_vec,
"target": target,
}
def train_model(
data_dir: str = "data",
model_dir: str = "models",
epochs: int = 50,
batch_size: int = 64,
learning_rate: float = 0.001,
device: str = "cuda",
) -> dict:
"""Train the form field classifier.
Returns:
Dictionary with training results and metrics.
"""
data_path = Path(data_dir)
model_path = Path(model_dir)
model_path.mkdir(parents=True, exist_ok=True)
# Load data
with open(data_path / "form_fields_train.json") as f:
train_data = json.load(f)
with open(data_path / "form_fields_test.json") as f:
test_data = json.load(f)
# Build vocabulary
vocab = Vocabulary()
vocab.build_from_samples(train_data)
vocab.save(str(model_path / "vocabulary.json"))
# Create datasets and loaders
train_dataset = FormFieldDataset(train_data, vocab)
test_dataset = FormFieldDataset(test_data, vocab)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# Create model
model = create_model(vocab_size=len(vocab.word2idx), device=device)
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# Training loop
print(f"\nTraining for {epochs} epochs on {device}...")
print("-" * 60)
best_accuracy = 0.0
train_start = time.time()
for epoch in range(epochs):
model.train()
total_loss = 0.0
correct = 0
total = 0
for batch in train_loader:
label_tokens = batch["label_tokens"].to(device)
name_tokens = batch["name_tokens"].to(device)
placeholder_tokens = batch["placeholder_tokens"].to(device)
input_type = batch["input_type"].to(device)
targets = batch["target"].to(device)
optimizer.zero_grad()
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type)
loss = criterion(logits, targets)
loss.backward()
optimizer.step()
total_loss += loss.item()
predictions = logits.argmax(dim=1)
correct += (predictions == targets).sum().item()
total += targets.size(0)
train_accuracy = correct / total
avg_loss = total_loss / len(train_loader)
# Evaluate on test set every 10 epochs
if (epoch + 1) % 10 == 0 or epoch == epochs - 1:
test_accuracy = evaluate(model, test_loader, device)
print(
f"Epoch {epoch + 1:3d}/{epochs} | "
f"Loss: {avg_loss:.4f} | "
f"Train Acc: {train_accuracy:.4f} | "
f"Test Acc: {test_accuracy:.4f}"
)
if test_accuracy > best_accuracy:
best_accuracy = test_accuracy
torch.save(model.state_dict(), model_path / "form_field_classifier_v1.pt")
train_duration = time.time() - train_start
print("-" * 60)
print(f"Training complete in {train_duration:.1f}s")
print(f"Best test accuracy: {best_accuracy:.4f}")
# Final detailed evaluation
print("\nPer-class accuracy:")
class_report = evaluate_detailed(model, test_loader, device)
for category, acc in class_report.items():
print(f" {category:20s}: {acc:.4f}")
results = {
"best_accuracy": best_accuracy,
"final_loss": avg_loss,
"epochs": epochs,
"training_time_seconds": train_duration,
"vocab_size": len(vocab.word2idx),
"model_path": str(model_path / "form_field_classifier_v1.pt"),
"class_report": class_report,
}
return results
def evaluate(model: FormFieldClassifier, loader: DataLoader, device: str) -> float:
"""Evaluate model accuracy on a data loader."""
model.eval()
correct = 0
total = 0
with torch.no_grad():
for batch in loader:
label_tokens = batch["label_tokens"].to(device)
name_tokens = batch["name_tokens"].to(device)
placeholder_tokens = batch["placeholder_tokens"].to(device)
input_type = batch["input_type"].to(device)
targets = batch["target"].to(device)
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type)
predictions = logits.argmax(dim=1)
correct += (predictions == targets).sum().item()
total += targets.size(0)
return correct / total
def evaluate_detailed(
model: FormFieldClassifier, loader: DataLoader, device: str
) -> dict[str, float]:
"""Per-class accuracy evaluation."""
model.eval()
class_correct: dict[int, int] = {i: 0 for i in range(len(FIELD_CATEGORIES))}
class_total: dict[int, int] = {i: 0 for i in range(len(FIELD_CATEGORIES))}
with torch.no_grad():
for batch in loader:
label_tokens = batch["label_tokens"].to(device)
name_tokens = batch["name_tokens"].to(device)
placeholder_tokens = batch["placeholder_tokens"].to(device)
input_type = batch["input_type"].to(device)
targets = batch["target"].to(device)
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type)
predictions = logits.argmax(dim=1)
for pred, target in zip(predictions, targets):
t = target.item()
class_total[t] += 1
if pred.item() == t:
class_correct[t] += 1
report = {}
for i, category in enumerate(FIELD_CATEGORIES):
if class_total[i] > 0:
report[category] = class_correct[i] / class_total[i]
else:
report[category] = 0.0
return report
def register_model(results: dict, dex_config: str = "dex.yaml"):
"""Register trained model in DataEngineX ModelRegistry."""
engine = DexEngine(dex_config)
registry = ModelRegistry()
artifact = ModelArtifact(
name="form_field_classifier",
version="1.0.0",
stage=ModelStage.STAGING,
metadata={
"framework": "pytorch",
"accuracy": results["best_accuracy"],
"training_time": results["training_time_seconds"],
"vocab_size": results["vocab_size"],
"num_classes": len(FIELD_CATEGORIES),
"model_path": results["model_path"],
},
)
registry.register(artifact, name="form_field_classifier", stage=ModelStage.STAGING)
print(f"Model registered: form_field_classifier v1.0.0 (STAGING)")
# Promote to production if accuracy meets threshold
if results["best_accuracy"] >= 0.85:
registry.register(artifact, name="form_field_classifier", stage=ModelStage.PRODUCTION)
print("Model promoted to PRODUCTION (accuracy >= 85%)")
return artifact
if __name__ == "__main__":
results = train_model()
register_model(results)Your model is trained, evaluated, and registered. Next up, you will expose it through a REST API so the future Chrome extension can call it to classify form fields in real time.
Register, Serve, and Document Gaps
Your model is trained, registered, and achieving 85%+ accuracy on unseen form fields. That is the hard part done.
Now you need to make the model accessible to other applications. A trained model sitting in a .pt file is useful to you, but useless to a browser extension that needs real-time predictions. In this step, you will wrap the model in a FastAPI REST endpoint so any HTTP client can classify form fields on demand.
In this step, get ready to:
- Verify the model registration in the DataEngineX ModelRegistry and note integration friction.
- Build a FastAPI prediction server with single and batch endpoints.
- File a GitHub Issue documenting integration gaps for future open-source contribution.
Review model registration and identify friction
When you ran train.py in the previous step, the register_model() function automatically registered your trained model in the DataEngineX ModelRegistry.
- Open src/formfield_classifier/train.py in VS Code and scroll to the register_model() function at the bottom of the file.
- Find the line where ModelArtifact is created.
artifact = ModelArtifact(
name="form_field_classifier",
version="1.0.0",
stage=ModelStage.STAGING,
metadata={
"framework": "pytorch",
"accuracy": results["best_accuracy"],
"training_time": results["training_time_seconds"],
"vocab_size": results["vocab_size"],
"num_classes": len(FIELD_CATEGORIES),
"model_path": results["model_path"],
},
)What did the registration do?
The ModelRegistry stored your model's metadata (accuracy, framework, path) in .dex/store.duckdb. It then promoted the model from STAGING to PRODUCTION because accuracy exceeded 85%.
Notice a friction point: the registration stores the model path as a string in metadata, but does not handle PyTorch serialization directly. The dataenginex.ml module ships a SklearnTrainer but has no equivalent PyTorchTrainer ABC. You handled serialization yourself with torch.save(), which works, but a first-class PyTorch integration would reduce boilerplate for future users.
Keep this friction in mind. You will file a GitHub Issue about it at the end of this step.
Build the FastAPI serving layer
The serving layer loads your trained model on startup and exposes HTTP endpoints that accept form field metadata and return predicted categories. This is the interface that a future Chrome extension will call.
- Create a new file src/formfield_classifier/serve.py in VS Code.
- Add the imports and Pydantic request/response models by pasting this code:
"""FastAPI serving layer for the form field classifier."""
import json
from pathlib import Path
from contextlib import asynccontextmanager
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from .dataset import FIELD_CATEGORIES, INPUT_TYPES, Vocabulary
from .model import FormFieldClassifier
class FieldInput(BaseModel):
"""Input schema for a single form field."""
label_text: str
input_type: str = "text"
name_attr: str = ""
placeholder: str = ""
class PredictionResult(BaseModel):
"""Output schema for a single prediction."""
predicted_category: str
confidence: float
top_3: list[dict]
class BatchInput(BaseModel):
"""Input schema for batch prediction."""
fields: list[FieldInput]
class BatchResult(BaseModel):
"""Output schema for batch prediction."""
predictions: list[PredictionResult]What does this code do?
The four Pydantic models define the API contract. FieldInput mirrors what the Chrome extension will send (label text, HTML input type, name attribute, placeholder). PredictionResult returns the predicted category with a confidence score and top 3 candidates.
BatchInput and BatchResult allow classifying an entire form's fields in a single HTTP request.
- Below the Pydantic models, add the app lifecycle and model loading code:
# Global model state
model: FormFieldClassifier | None = None
vocab: Vocabulary | None = None
device: str = "cuda" if torch.cuda.is_available() else "cpu"
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Load model on startup."""
global model, vocab
model_dir = Path("models")
model_path = model_dir / "form_field_classifier_v1.pt"
vocab_path = model_dir / "vocabulary.json"
if not model_path.exists():
raise RuntimeError(f"Model not found at {model_path}. Train the model first.")
# Load vocabulary
vocab = Vocabulary()
vocab.load(str(vocab_path))
# Load model
model = FormFieldClassifier(vocab_size=len(vocab.word2idx))
model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
model.to(device)
model.eval()
print(f"Model loaded on {device} | Vocab size: {len(vocab.word2idx)}")
yield
print("Shutting down...")
app = FastAPI(
title="Form Field Classifier API",
description="Classifies job application form fields from HTML metadata",
version="1.0.0",
lifespan=lifespan,
)What does this code do?
The lifespan async context manager runs once when the server starts. It loads the vocabulary and model weights from disk, moves the model to GPU (or CPU as fallback), and sets it to evaluation mode with model.eval().
Using a lifespan handler means the model loads once at startup instead of on every request. The weights_only=True flag in torch.load() prevents arbitrary code execution from untrusted model files.
- Add the prediction logic and all three endpoints below the app instance:
def predict_field(field: FieldInput) -> PredictionResult:
"""Run prediction for a single field."""
if model is None or vocab is None:
raise RuntimeError("Model not loaded")
# Encode inputs
label_tokens = torch.tensor(
[vocab.encode(field.label_text)], dtype=torch.long, device=device
)
name_tokens = torch.tensor(
[vocab.encode(field.name_attr)], dtype=torch.long, device=device
)
placeholder_tokens = torch.tensor(
[vocab.encode(field.placeholder)], dtype=torch.long, device=device
)
input_type_vec = torch.zeros(1, len(INPUT_TYPES), dtype=torch.float, device=device)
if field.input_type in INPUT_TYPES:
type_idx = INPUT_TYPES.index(field.input_type)
input_type_vec[0, type_idx] = 1.0
# Predict
with torch.no_grad():
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type_vec)
probabilities = torch.softmax(logits, dim=1)[0]
# Get top predictions
top_values, top_indices = probabilities.topk(3)
top_3 = [
{"category": FIELD_CATEGORIES[idx.item()], "confidence": round(val.item(), 4)}
for val, idx in zip(top_values, top_indices)
]
predicted_idx = top_indices[0].item()
return PredictionResult(
predicted_category=FIELD_CATEGORIES[predicted_idx],
confidence=round(top_values[0].item(), 4),
top_3=top_3,
)
@app.get("/health")
async def health():
"""Health check endpoint."""
return {
"status": "healthy",
"model_loaded": model is not None,
"device": device,
}
@app.post("/predict", response_model=PredictionResult)
async def predict(field: FieldInput):
"""Predict the category of a single form field."""
try:
return predict_field(field)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/batch-predict", response_model=BatchResult)
async def batch_predict(batch: BatchInput):
"""Predict categories for multiple form fields."""
try:
predictions = [predict_field(field) for field in batch.fields]
return BatchResult(predictions=predictions)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)What does this code do?
predict_field() encodes the incoming field metadata into tensors, runs a forward pass through the model (inside torch.no_grad() for inference efficiency), and returns the top 3 predicted categories with confidence scores.
The three endpoints cover different use cases: /health for monitoring, /predict for classifying a single field, and /batch-predict for classifying an entire form at once.
- Save src/formfield_classifier/serve.py.
✔️ Awesome, I've got everything!
Great. Double check you have saved the file before starting the server.
ⓧ I'd like to double check the full code
"""FastAPI serving layer for the form field classifier."""
import json
from pathlib import Path
from contextlib import asynccontextmanager
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from .dataset import FIELD_CATEGORIES, INPUT_TYPES, Vocabulary
from .model import FormFieldClassifier
class FieldInput(BaseModel):
"""Input schema for a single form field."""
label_text: str
input_type: str = "text"
name_attr: str = ""
placeholder: str = ""
class PredictionResult(BaseModel):
"""Output schema for a single prediction."""
predicted_category: str
confidence: float
top_3: list[dict]
class BatchInput(BaseModel):
"""Input schema for batch prediction."""
fields: list[FieldInput]
class BatchResult(BaseModel):
"""Output schema for batch prediction."""
predictions: list[PredictionResult]
# Global model state
model: FormFieldClassifier | None = None
vocab: Vocabulary | None = None
device: str = "cuda" if torch.cuda.is_available() else "cpu"
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Load model on startup."""
global model, vocab
model_dir = Path("models")
model_path = model_dir / "form_field_classifier_v1.pt"
vocab_path = model_dir / "vocabulary.json"
if not model_path.exists():
raise RuntimeError(f"Model not found at {model_path}. Train the model first.")
# Load vocabulary
vocab = Vocabulary()
vocab.load(str(vocab_path))
# Load model
model = FormFieldClassifier(vocab_size=len(vocab.word2idx))
model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
model.to(device)
model.eval()
print(f"Model loaded on {device} | Vocab size: {len(vocab.word2idx)}")
yield
print("Shutting down...")
app = FastAPI(
title="Form Field Classifier API",
description="Classifies job application form fields from HTML metadata",
version="1.0.0",
lifespan=lifespan,
)
def predict_field(field: FieldInput) -> PredictionResult:
"""Run prediction for a single field."""
if model is None or vocab is None:
raise RuntimeError("Model not loaded")
# Encode inputs
label_tokens = torch.tensor(
[vocab.encode(field.label_text)], dtype=torch.long, device=device
)
name_tokens = torch.tensor(
[vocab.encode(field.name_attr)], dtype=torch.long, device=device
)
placeholder_tokens = torch.tensor(
[vocab.encode(field.placeholder)], dtype=torch.long, device=device
)
input_type_vec = torch.zeros(1, len(INPUT_TYPES), dtype=torch.float, device=device)
if field.input_type in INPUT_TYPES:
type_idx = INPUT_TYPES.index(field.input_type)
input_type_vec[0, type_idx] = 1.0
# Predict
with torch.no_grad():
logits = model(label_tokens, name_tokens, placeholder_tokens, input_type_vec)
probabilities = torch.softmax(logits, dim=1)[0]
# Get top predictions
top_values, top_indices = probabilities.topk(3)
top_3 = [
{"category": FIELD_CATEGORIES[idx.item()], "confidence": round(val.item(), 4)}
for val, idx in zip(top_values, top_indices)
]
predicted_idx = top_indices[0].item()
return PredictionResult(
predicted_category=FIELD_CATEGORIES[predicted_idx],
confidence=round(top_values[0].item(), 4),
top_3=top_3,
)
@app.get("/health")
async def health():
"""Health check endpoint."""
return {
"status": "healthy",
"model_loaded": model is not None,
"device": device,
}
@app.post("/predict", response_model=PredictionResult)
async def predict(field: FieldInput):
"""Predict the category of a single form field."""
try:
return predict_field(field)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/batch-predict", response_model=BatchResult)
async def batch_predict(batch: BatchInput):
"""Predict categories for multiple form fields."""
try:
predictions = [predict_field(field) for field in batch.fields]
return BatchResult(predictions=predictions)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)Now start the server. The --app-dir flag tells uvicorn to add src to the Python path so it can find the formfield_classifier package.
- Start the prediction server by running this command from the formfield-classifier folder:
uvicorn formfield_classifier.serve:app --app-dir src --host 0.0.0.0 --port 8000What does this command do?
uvicorn is the ASGI server that runs your FastAPI app. formfield_classifier.serve:app tells it to import the app object from the formfield_classifier.serve module. --app-dir src adds src/ to the PYTHONPATH so Python can resolve the package import.
- You should see output confirming the model loaded and the server is running at http://0.0.0.0:8000.
- Open a second terminal and test the /predict endpoint by running:
curl -X POST http://localhost:8000/predict \
-H "Content-Type: application/json" \
-d '{"label_text": "Email Address", "input_type": "email", "name_attr": "email", "placeholder": "you@example.com"}'- You should see a JSON response with "predicted_category": "email" and a confidence value above 0.8.
- Test the /health endpoint by running:
curl http://localhost:8000/health- You should see {"status":"healthy","model_loaded":true,"device":"cuda"}.
Server not starting?
- If you see a ModuleNotFoundError, make sure you are running the uvicorn command from the formfield-classifier project root (not from inside src/).
- If you see a RuntimeError about the model not being found, confirm that models/form_field_classifier_v1.pt exists. Training must have completed in the previous step.
- If port 8000 is in use, check for other running processes and stop them first.
Help me debug why my FastAPI server won't start for the form field classifier.
File a GitHub Issue documenting integration gaps
During this project, you discovered that dataenginex.ml provides a SklearnTrainer but no equivalent abstraction for PyTorch workflows. This is a real integration gap that other users would benefit from knowing about.
- Open https://github.com/TheDataEngineX/dex in your browser.
- Click the **Issues** tab near the top of the repository page.
- Click **New issue**.
- Set the title to something descriptive, for example: Add PyTorchTrainer ABC alongside SklearnTrainer.
- In the issue body, include three sections:
- Add the labels enhancement and ml to your issue.
- Click **Submit new issue**.
Why file an Issue instead of a PR?
Good open-source contribution starts with communication. Filing an Issue with a clear proposal lets maintainers weigh in on API design before you write code. It also leaves a trail for future contributors who hit the same friction.
Other gaps worth documenting
You may also want to mention path-based model artifact registration. The current register() method works, but a helper that accepts a file path and handles serialization metadata (framework, input shape, required vocab) would reduce boilerplate for non-sklearn models.
Secret mission
Fine-tune DistilBERT for Semantic Field Matching
Your from-scratch model handles common fields like "First Name" and "Email" with high confidence. But what about ambiguous labels like "Additional Information", "Tell us about yourself", or "Anything else?" These need deeper semantic understanding that a bag-of-embeddings model cannot provide. In this secret mission, you will fine-tune a pretrained DistilBERT model and build an ensemble predictor that combines both models for robust field classification.
Clean Up Your Resources
Clean Up Your Resources
This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.
Resources you used:
- The formfield-classifier/ project directory (source code, training data, model weights, DuckDB store).
- FastAPI server process (if currently running on port 8000).
- Python packages installed via pip (torch, dataenginex, fastapi, uvicorn, transformers).
Keep everything running
No action needed. Choose this if you plan to build the Chrome extension that calls this API, or if you want to continue experimenting with the model.
- Your trained models persist in models/ and are registered in .dex/store.duckdb.
- Start the prediction server anytime by running:
- python -m formfield_classifier.serve from inside the formfield-classifier/ directory.
- The installed Python packages (torch, dataenginex, fastapi, uvicorn, transformers) remain available for other projects.
Pause - I'll come back to this later
Stop any running processes to free up memory. Your project files and model weights stay on disk.
- If the FastAPI server is running, press Ctrl+C in the terminal to stop it.
- All files in formfield-classifier/ remain on disk, including your trained model weights and the DuckDB store.
- When you return, open a terminal in the formfield-classifier/ directory and run the server again.
Delete - I don't want to use this again
Remove all project files and start fresh. This deletes your trained models, training data, and the DuckDB store permanently.
- If the FastAPI server is running, press Ctrl+C in the terminal to stop it.
- Delete the project directory by running:
cd ~/Desktop
rm -rf formfield-classifierOn Windows, run this instead:
cd ~\Desktop
rmdir /s /q formfield-classifier- Optionally uninstall the Python packages by running:
pip uninstall torch torchvision torchaudio dataenginex fastapi uvicorn transformers -y- Confirm the directory is removed by running ls ~/Desktop (macOS/Linux) or dir ~\Desktop (Windows). You should no longer see formfield-classifier in the output.
That's a Wrap!
That's a Wrap!
Nice work! You just built a production-ready ML pipeline from the ground up. Starting from raw form field metadata, you designed a neural network architecture, trained it on your GPU, and shipped a live prediction API.
You've learned how to:
- Design and train a custom PyTorch neural network from scratch that classifies 18 form field categories with GPU-accelerated training on your NVIDIA Quadro T2000.
- Build an end-to-end MLOps pipeline covering synthetic dataset generation, model registration in DataEngineX ModelRegistry with stage management, and prediction serving via a FastAPI REST endpoint.
- Contribute to an open-source ecosystem by filing GitHub Issues on TheDataEngineX/dex with proposed backward-compatible solutions for PyTorch integration gaps.
- Secret Mission: Fine-tuned DistilBERT for semantic understanding of ambiguous form fields and built an ensemble predictor combining both models for robust classification.
Ready to quiz yourself?