Scope: this document covers the model-training portion of notebooks/PR_CLassifier.ipynb, starting at the cell labeled "Model Training starts from here" and ending after validation/evaluation. It focuses on the executed DistilBERT multi-task trainer, not the earlier data preparation phase.
The notebook trains a single shared encoder with two task-specific heads:
- Sentiment head: 3-way multi-class classification for
Negative,Neutral,Positive - GitHub label head: 5-way multi-label classification for
bug,enhancement,documentation,test,request
The key design choice is that the model learns a shared semantic representation from text, then routes that representation into the correct prediction head depending on the task being optimized.
A classical baseline is also evaluated in the notebook for context, but the rest of this document focuses on the DistilBERT-based multi-task model that becomes the production checkpoint.
💡 NOTE: The sentiment task is trained on the
cardiffnlp/tweet_evalsentiment corpus, while the GitHub label task is trained on the cleaned GitHub issue/PR corpus. That separation gives the encoder broad language coverage without forcing the sentiment labels to come from the same domain as the GitHub taxonomy.
The implementation uses DistilBertModel as the backbone and then adds a lightweight projection block plus two classification heads. The forward path is intentionally simple:
- Tokenize the input text
- Encode with DistilBERT
- Mean-pool the token states using the attention mask
- Pass the pooled vector through a shared projection layer
- Branch into the sentiment head or the GitHub label head based on the
taskflag
flowchart TD
A[Input text\nTitle + Body or Body-only text] --> B[DistilBERT tokenizer]
B --> C[Input IDs]
B --> D[Attention mask]
C --> E[DistilBERT encoder\noutputs last hidden state]
D --> E
E --> F[Masked mean pooling]
F --> G[Shared pre-classifier\nLinear 768 -> 768]
G --> H[ReLU]
H --> I[Dropout p=0.2]
I --> J{Task routing}
J -->|sentiment| K[Sentiment head\nLinear 768 -> 3]
J -->|github| L[GitHub head\nLinear 768 -> 5]
K --> M[CrossEntropyLoss]
L --> N[BCEWithLogitsLoss\nwith pos_weight]
M --> O[Backprop]
N --> O
O --> P[AdamW update]
class MultiTaskDistilBERT(nn.Module):
def __init__(self, num_sentiment=3, num_github=5):
super().__init__()
self.distilbert = DistilBertModel.from_pretrained('distilbert-base-uncased')
self.pre_classifier = nn.Linear(768, 768) # shared projection before branching
self.dropout = nn.Dropout(0.2) # regularize the shared latent space
self.sentiment_head = nn.Linear(768, num_sentiment) # 3-class sentiment logits
self.github_head = nn.Linear(768, num_github) # 5 independent label logits
def forward(self, input_ids, attention_mask, task):
outputs = self.distilbert(input_ids=input_ids, attention_mask=attention_mask)
hidden_state = outputs[0] # shape: [batch, seq_len, 768]
# Masked mean pooling keeps padding tokens out of the sentence representation.
mask = attention_mask.unsqueeze(-1).expand(hidden_state.size()).float()
mean_pooled = torch.sum(hidden_state * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)
x = nn.ReLU()(self.pre_classifier(mean_pooled))
x = self.dropout(x)
return self.sentiment_head(x) if task == 'sentiment' else self.github_head(x)DistilBERT does not expose a classic BERT-style pooler in the same way some encoders do, so the notebook uses attention-mask-aware mean pooling instead of relying on a single special token. That is a strong choice for issue/PR text because:
- the representation uses information from the full sequence, not just one position
- padding does not contaminate the hidden state average
- long bodies tend to distribute semantic signal across multiple spans, so averaging is more stable than a raw first-token readout
Mathematically, the pooled vector is:
where
💡 NOTE: The
taskargument is a simple runtime switch. In this notebook, any call withtask='sentiment'routes to the 3-way head, and every other task value routes to the GitHub head. That keeps the training loop explicit, but the caller must pass the correct task string.
The notebook deliberately uses two different losses because the prediction geometries are different.
The sentiment task is mutually exclusive. For each example, exactly one class should win among 3 categories. That makes it a standard multi-class classification problem.
CrossEntropyLoss is the right choice because it combines a softmax over the logits with the negative log-likelihood of the target class. The model outputs three raw logits, and the loss pushes probability mass toward the correct class.
The GitHub task is multi-label, not multi-class. A single issue or PR can be both bug and test, or enhancement and request, at the same time. Each label is therefore modeled as an independent Bernoulli target.
BCEWithLogitsLoss is the right choice because it:
- applies a sigmoid internally in a numerically stable way
- computes binary cross-entropy per label
- allows each label to activate independently
The notebook also supplies pos_weight to reweight rare labels. That is important because the GitHub taxonomy is strongly imbalanced.
The training loop computes one loss for the GitHub batch and one loss for the sentiment batch, then adds them before backpropagation:
This is an equal-weight multi-task objective. There is no additional task coefficient, so both heads contribute symmetrically to the update step.
github_counts = github_train_df[target_cols].sum(axis=0).values
total_train = len(github_train_df)
pos_weights = torch.tensor((total_train - github_counts) / github_counts, dtype=torch.float)
optimizer = AdamW(model.parameters(), lr=2e-5) # AdamW with decoupled weight decay
scaler = torch.cuda.amp.GradScaler() # stabilize mixed-precision training
loss_fn_github = nn.BCEWithLogitsLoss(pos_weight=pos_weights.to(device))
loss_fn_twit = nn.CrossEntropyLoss()
⚠️ WARNING: The notebook does not use focal loss or class-specific sampling for the GitHub head. Thepos_weightcorrection helps with imbalance, but very rare labels can still remain threshold-sensitive at inference time.
The notebook uses AdamW with a learning rate of 2e-5.
AdamW is a strong fit here because it combines Adam-style adaptive moments with decoupled weight decay, which tends to generalize better than classic L2 regularization when fine-tuning transformers.
Because the notebook does not override weight_decay, the optimizer runs with the PyTorch default decoupled weight decay behavior. That is a reasonable choice for transformer fine-tuning and acts as a mild regularizer on the shared encoder and both heads.
Training uses torch.cuda.amp.autocast() plus GradScaler.
That matters because:
- it reduces GPU memory consumption
- it speeds up forward and backward passes on CUDA hardware
- it keeps the large transformer backbone tractable under batch size 16
The notebook imports get_linear_schedule_with_warmup, but the executed training loop does not instantiate or step a scheduler. In other words, the run uses a constant learning rate throughout training.
⚠️ WARNING: There is no warmup or decay schedule in the executed training cells. If you later increase sequence length, batch size, or backbone size, adding a scheduler would be one of the first stabilization upgrades to consider.
The notebook trains the two tasks in lockstep.
- The GitHub loader is the longer stream
- The Twitter loader is cycled with
itertools.cycle(...) - Each iteration consumes one GitHub batch and one sentiment batch
- The optimizer receives the sum of both task losses once per step
This is a practical multi-task schedule because it prevents the shorter sentiment dataset from exhausting early and silently shortening the epoch.
sequenceDiagram
participant GH as GitHub loader
participant TW as Twitter loader
participant M as MultiTaskDistilBERT
participant S as GradScaler
participant O as AdamW
loop each epoch
O->>O: zero_grad()
GH->>M: batch(task="github")
M-->>GH: github logits
GH->>GH: BCEWithLogitsLoss
TW->>M: batch(task="sentiment")
M-->>TW: sentiment logits
TW->>TW: CrossEntropyLoss
GH->>S: scale(loss_g + loss_t).backward()
S->>O: step()
S->>S: update()
end
for git_batch, twit_batch in tqdm.tqdm(zip(github_loader, cycle(twitter_loader)), total=len(github_loader)):
optimizer.zero_grad()
# GitHub multi-label update
with torch.cuda.amp.autocast():
g_out = model(git_batch['input_ids'].to(device), git_batch['attention_mask'].to(device), 'github')
loss_g = loss_fn_github(g_out, git_batch['labels'].to(device))
# Sentiment multi-class update
with torch.cuda.amp.autocast():
t_out = model(twit_batch['input_ids'].to(device), twit_batch['attention_mask'].to(device), 'sentiment')
loss_t = loss_fn_twit(t_out, twit_batch['labels'].to(device))
# One scalar objective, two tasks, one backward pass
scaler.scale(loss_g + loss_t).backward()
scaler.step(optimizer)
scaler.update()Because the model shares the DistilBERT backbone and the pre_classifier, gradients from both tasks accumulate in the shared layers. That produces a genuine multi-task effect:
- GitHub batches teach the encoder how issue/PR language maps to label semantics
- sentiment batches teach the same encoder how tone and polarity manifest in noisy short-form text
- the shared projection layer learns a compact latent space that is useful for both tasks
This is the central engineering idea in the notebook.
💡 NOTE: The notebook uses the sum of validation macro F1 scores to pick the best checkpoint, not the raw training loss. That means model selection is aligned with the downstream metric that actually matters for the two heads.
Evaluation is done with model.eval() and torch.no_grad(). The notebook uses task-appropriate decoding rules:
- Sentiment:
argmax(logits, dim=1) - GitHub labels:
sigmoid(logits) > 0.5
That thresholding choice is standard for multi-label heads because each label is independent. The 0.5 cutoff is easy to interpret and works well as a default before probability calibration.
- Sentiment is a single-label classification problem, so accuracy and class-wise F1 are informative.
- GitHub labels are multi-label and imbalanced, so micro F1, macro F1, and samples F1 tell a more complete story.
Macro F1 matters especially here because it weights each label equally. That exposes minority-label behavior that micro F1 can hide.
def evaluate(model, loader, task):
model.eval()
all_preds, all_labels = [], []
with torch.no_grad():
for batch in loader:
ids, mask, labels = batch['input_ids'].to(device), batch['attention_mask'].to(device), batch['labels'].to(device)
outputs = model(ids, mask, task)
if task == 'sentiment':
preds = torch.argmax(outputs, dim=1).cpu().numpy()
else:
preds = (torch.sigmoid(outputs) > 0.5).cpu().numpy()
all_preds.extend(preds)
all_labels.extend(labels.cpu().numpy())
f1 = f1_score(all_labels, all_preds, average='macro')
return f1The final report in the notebook is generated on the validation loaders, not on the held-out test loaders. So the numbers below should be treated as validation performance used for checkpoint selection and model comparison.
⚠️ WARNING: A validation report is not the same thing as a blind final test. If you later want a publishable benchmark, run the saved checkpoint on the untouched test splits as a separate final pass.
| Hyperparameter | Value | Why it matters |
|---|---|---|
| Backbone | DistilBERT base uncased | Compact transformer with 768-d hidden states |
| Pooling strategy | Attention-mask mean pooling | Uses the full sequence without padding leakage |
| Shared projection | Linear 768 -> 768 | Learns a task-agnostic latent bottleneck |
| Dropout | 0.2 | Regularizes the shared representation |
| Sentiment head | Linear 768 -> 3 | Multi-class sentiment logits |
| GitHub head | Linear 768 -> 5 | Independent logits for each label |
| Max sequence length | 256 | Balances context coverage and memory usage |
| Batch size | 16 | Notebook comment notes T4 stability |
| Optimizer | AdamW | Good default for transformer fine-tuning |
| Learning rate | 2e-5 | Conservative fine-tuning step size |
| Loss - sentiment | CrossEntropyLoss | Correct objective for mutually exclusive classes |
| Loss - GitHub | BCEWithLogitsLoss | Correct objective for independent labels |
| Positive weighting | (N_neg / N_pos) per label | Compensates for label imbalance |
| AMP | autocast + GradScaler | Faster and more memory-efficient CUDA training |
| Scheduler | None executed | Imported, but no warmup/decay was stepped in the notebook |
| Total epochs | 4 | 2 initial + 2 continuation epochs |
| Checkpoint selection | Sum of validation macro F1 scores | Balances both tasks equally |
| Component | Setting | Practical impact |
|---|---|---|
| Device selection | cuda if available else cpu |
Runs on GPU when available, otherwise falls back safely |
| Precision mode | Mixed precision on CUDA | Reduces memory footprint and accelerates training |
| Update pattern | Paired GitHub + sentiment batches | Maintains one multi-task update stream per step |
| Epoch length | len(github_loader) steps |
About 3,253 paired updates per epoch |
| Model artifact | ../models/pr_classifier.pt |
Final production checkpoint used by the backend |
| Serving model mirror | src/models/model_arch.py |
Same architecture is reloaded in the API layer |
The notebook runs the model for four epochs total and reports validation macro F1 for both tasks after each epoch.
| Epoch | GitHub Macro F1 | Twitter Macro F1 | Checkpoint Status |
|---|---|---|---|
| 1 | 0.6219 | 0.7070 | Initial improvement |
| 2 | 0.6574 | 0.7103 | Better joint score |
| 3 | 0.6849 | 0.7058 | Best checkpoint saved |
| 4 | 0.6722 | 0.6927 | Slight regression |
The best joint checkpoint is written during the extended training block to ../models/pr_classifier.pt. That is the artifact consumed by the FastAPI backend on startup.
💡 NOTE: The validation trace shows the GitHub head continuing to improve through epoch 3 while the sentiment head plateaus slightly earlier. That is a normal multi-task pattern when one task is harder or more imbalanced than the other.
The notebook’s final evaluation uses the validation splits and prints a full classification report for both tasks.
| Task | Metric | Value | Interpretation |
|---|---|---|---|
| Sentiment | Accuracy | 0.71 | Strong 3-way classification score |
| Sentiment | Macro F1 | 0.69 | Balanced across sentiment classes |
| Sentiment | Weighted F1 | 0.71 | Support-weighted summary |
| GitHub labels | Micro F1 | 0.75 | Global multi-label performance |
| GitHub labels | Macro F1 | 0.67 | Minority-label sensitive view |
| GitHub labels | Weighted F1 | 0.78 | Support-weighted multi-label score |
| GitHub labels | Samples F1 | 0.77 | Example-level label quality |
| Label | Precision | Recall | F1 | Support |
|---|---|---|---|---|
| bug | 0.87 | 0.87 | 0.87 | 3610 |
| enhancement | 0.77 | 0.81 | 0.79 | 2117 |
| documentation | 0.33 | 0.86 | 0.47 | 535 |
| test | 0.37 | 0.87 | 0.52 | 516 |
| request | 0.61 | 0.86 | 0.71 | 879 |
The rare labels are intentionally recall-heavy. That is a direct consequence of the imbalance correction and the 0.5 sigmoid threshold. For a production labeling pipeline, that tradeoff is often preferable to missing a valid label entirely, especially for automation workflows that can be corrected by humans later.
| Class | Precision | Recall | F1 | Support |
|---|---|---|---|---|
| Negative | 0.60 | 0.66 | 0.63 | 312 |
| Neutral | 0.70 | 0.67 | 0.68 | 869 |
| Positive | 0.76 | 0.77 | 0.76 | 819 |
The training notebook does not end as an isolated experiment. Its final checkpoint is loaded by the API layer at runtime:
- src/api/backend.py loads ../models/pr_classifier.pt
- src/models/model_arch.py defines the same
MultiTaskDistilBERTarchitecture used during training - the Flask app and n8n workflow then call that backend to label live GitHub issues and PRs automatically
This is why the training design matters: the model is not just scoring a notebook cell, it is the decision engine behind the live automation pipeline.
The training phase is a compact but well-structured multi-task fine-tuning pipeline:
- DistilBERT provides the shared semantic backbone
- masked mean pooling creates a robust sequence representation
- a shared projection layer feeds two task-specific heads
CrossEntropyLosshandles sentimentBCEWithLogitsLosswith positive weighting handles GitHub labels- AdamW plus AMP keeps the run efficient on commodity GPU hardware
- validation macro F1 drives checkpoint selection
The result is a production-friendly dual-head classifier that can support both sentiment analysis and automatic GitHub label suggestion in the same automation stack.