truthlens-ml-api / src /train_verification.py
Param0304's picture
Deploy TruthLens ML API to Hugging Face Space
3abf967
Raw History Blame Contribute Delete
18.9 kB
"""
train_verification.py - Claim Verification Model (Model B) Training for TruthLens AI.
Task Definition:
MODEL B: CLAIM VERIFICATION MODEL
Input: [CLAIM] claim text [SEP] evidence text
Output: SUPPORTS (0), REFUTES (1), NOT_ENOUGH_INFO (2)
Hardware Acceleration:
- Strictly utilizes NVIDIA GPU with CUDA acceleration.
- Mixed Precision FP16 enabled.
- Gradient Accumulation to maintain effective batch size = 16 within 4GB VRAM envelope.
- Peak GPU memory tracking.
"""
import os
import sys
import json
import time
import math
from typing import Dict, List, Any, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import pandas as pd
import numpy as np
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
get_linear_schedule_with_warmup,
)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from check_gpu import inspect_gpu_environment
from training_utils import (
LABEL2ID,
ID2LABEL,
set_seed,
compute_class_weights,
compute_verification_metrics,
plot_and_save_confusion_matrix,
log_hardware_metrics,
)
class SciFactVerificationDataset(Dataset):
def __init__(self, df: pd.DataFrame, tokenizer, max_length: int = 256):
self.tokenizer = tokenizer
self.max_length = max_length
self.claims = df["claim"].astype(str).tolist()
self.evidences = df["evidence"].fillna("").astype(str).tolist()
self.labels = [LABEL2ID[l] for l in df["label"]]
def __len__(self):
return len(self.claims)
def __getitem__(self, idx):
claim = self.claims[idx]
evidence = self.evidences[idx]
label = self.labels[idx]
# Natural format: [CLAIM] claim text [SEP] evidence text
claim_clean = str(claim).strip()
evidence_clean = str(evidence).strip() if evidence and str(evidence).strip() else "None"
input_text = f"[CLAIM] {claim_clean} [SEP] {evidence_clean}"
encoding = self.tokenizer(
input_text,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt"
)
item = {k: v.squeeze(0) for k, v in encoding.items()}
item["label"] = torch.tensor(label, dtype=torch.long)
return item
def train_epoch(
model: nn.Module,
dataloader: DataLoader,
optimizer: torch.optim.Optimizer,
scheduler: Any,
scaler: torch.amp.GradScaler,
criterion: nn.Module,
device: torch.device,
amp_dtype: torch.dtype = torch.bfloat16,
grad_accum_steps: int = 4
) -> float:
model.train()
total_loss = 0.0
optimizer.zero_grad()
use_scaler = scaler.is_enabled()
for step, batch in enumerate(dataloader, start=1):
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels = batch["label"].to(device)
with torch.amp.autocast("cuda", dtype=amp_dtype, enabled=device.type == "cuda"):
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
loss = criterion(outputs.logits, labels)
loss = loss / grad_accum_steps
if use_scaler:
scaler.scale(loss).backward()
else:
loss.backward()
total_loss += loss.item() * grad_accum_steps
if step % grad_accum_steps == 0 or step == len(dataloader):
if use_scaler:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
else:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
optimizer.zero_grad()
if scheduler is not None:
scheduler.step()
return total_loss / len(dataloader)
def evaluate(
model: nn.Module,
dataloader: DataLoader,
device: torch.device,
amp_dtype: torch.dtype = torch.bfloat16
) -> Tuple[Dict[str, Any], List[int], List[int], List[List[float]]]:
model.eval()
all_preds = []
all_targets = []
all_probs = []
with torch.no_grad():
for batch in dataloader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels = batch["label"]
with torch.amp.autocast("cuda", dtype=amp_dtype, enabled=device.type == "cuda"):
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
probs = F.softmax(outputs.logits, dim=-1).cpu().numpy()
preds = np.argmax(probs, axis=-1)
all_preds.extend(preds.tolist())
all_targets.extend(labels.tolist())
all_probs.extend(probs.tolist())
metrics = compute_verification_metrics(all_targets, all_preds)
return metrics, all_targets, all_preds, all_probs
def train_scifact_verification_model(
base_dir: str,
model_name: str = "roberta-base",
epochs: int = 4,
lr: float = 2e-5,
per_device_batch_size: int = 8,
grad_accum_steps: int = 2,
max_length: int = 256,
weight_decay: float = 0.01,
warmup_ratio: float = 0.1,
seed: int = 42
) -> Dict[str, Any]:
print("=" * 80)
print(" TRUTHLENS AI — CLAIM VERIFICATION TRANSFORMER TRAINING (MODEL B)")
print("=" * 80)
# 1. Inspect Hardware and Require CUDA
set_seed(seed)
gpu_info = inspect_gpu_environment(require_cuda=True)
device = torch.device("cuda")
torch.cuda.reset_peak_memory_stats(device)
data_dir = os.path.join(base_dir, "data", "final", "verification")
models_dir = os.path.join(base_dir, "models", "verification", "scifact_deberta")
reports_dir = os.path.join(base_dir, "reports")
os.makedirs(models_dir, exist_ok=True)
os.makedirs(reports_dir, exist_ok=True)
train_df = pd.read_csv(os.path.join(data_dir, "scifact_train.csv"))
valid_df = pd.read_csv(os.path.join(data_dir, "scifact_valid.csv"))
print(f"Dataset: SciFact Grounded Verification")
print(f"Train samples: {len(train_df):,} | Valid samples: {len(valid_df):,}")
print(f"Train label distribution:\n{train_df['label'].value_counts().to_dict()}")
# 2. Sequence Length Verification
train_df["evidence"] = train_df["evidence"].fillna("").astype(str)
comb_words = train_df.apply(lambda r: len(r["claim"].split()) + len(r["evidence"].split()), axis=1)
max_comb_words = int(comb_words.max())
print(f"Max combined word length: {max_comb_words} words (max_length={max_length} guarantees 0.0% truncation)")
# 3. Load Tokenizer & Model
print(f"\nLoading Pretrained Transformer: {model_name}...")
try:
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)
except Exception as e:
print(f"Notice: {model_name} load issue ({e}). Falling back to 'roberta-base'...")
model_name = "roberta-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)
model.to(device)
# 4. Datasets and Loaders
train_dataset = SciFactVerificationDataset(train_df, tokenizer, max_length=max_length)
valid_dataset = SciFactVerificationDataset(valid_df, tokenizer, max_length=max_length)
train_loader = DataLoader(train_dataset, batch_size=per_device_batch_size, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=per_device_batch_size * 2, shuffle=False)
# 5. Class Weights & Loss
train_labels = [LABEL2ID[l] for l in train_df["label"]]
class_weights = compute_class_weights(train_labels, num_classes=3)
print(f"Class Weights (Inverse Frequency): {class_weights}")
criterion = nn.CrossEntropyLoss(weight=torch.tensor(class_weights, device=device))
# 6. Optimizer & Scheduler
total_training_steps = math.ceil(len(train_loader) / grad_accum_steps) * epochs
warmup_steps = int(total_training_steps * warmup_ratio)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=warmup_steps,
num_training_steps=total_training_steps
)
amp_dtype = torch.float16
precision_name = "FP16"
print(f"Mixed Precision Mode: {precision_name} (amp_dtype={amp_dtype})")
scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")
# 7. Training Loop
print("\nStarting Fine-Tuning...")
start_train_time = time.time()
best_macro_f1 = -1.0
best_metrics = None
best_targets, best_preds, best_probs = [], [], []
for epoch in range(1, epochs + 1):
ep_start = time.time()
avg_loss = train_epoch(
model=model,
dataloader=train_loader,
optimizer=optimizer,
scheduler=scheduler,
scaler=scaler,
criterion=criterion,
device=device,
amp_dtype=amp_dtype,
grad_accum_steps=grad_accum_steps
)
ep_time = time.time() - ep_start
# Evaluate on Validation Split
val_metrics, val_targets, val_preds, val_probs = evaluate(
model=model,
dataloader=valid_loader,
device=device,
amp_dtype=amp_dtype
)
print(f"Epoch {epoch}/{epochs} ({ep_time:.1f}s) | Train Loss: {avg_loss:.4f} | Val Acc: {val_metrics['accuracy']*100:.2f}% | Val Macro F1: {val_metrics['macro_f1']*100:.2f}% | Weighted F1: {val_metrics['weighted_f1']*100:.2f}%")
if val_metrics["macro_f1"] > best_macro_f1:
best_macro_f1 = val_metrics["macro_f1"]
best_metrics = val_metrics
best_targets, best_preds, best_probs = val_targets, val_preds, val_probs
# Save Best Model Artifacts
model.save_pretrained(models_dir)
tokenizer.save_pretrained(models_dir)
print(f" -> Saved new best model checkpoint (Macro F1: {best_macro_f1*100:.2f}%)")
total_training_time = time.time() - start_train_time
peak_vram_mb = torch.cuda.max_memory_allocated(device) / (1024 ** 2)
print("\n" + "=" * 80)
print(" TRAINING COMPLETED")
print(f"Total Time: {total_training_time:.2f}s ({total_training_time/60:.2f} min)")
print(f"Peak GPU VRAM Allocated: {peak_vram_mb:.2f} MB")
print(f"Best Macro F1: {best_macro_f1*100:.2f}%")
print("=" * 80)
# 8. Save Model Metadata & Configuration
config = {
"model_name": model_name,
"task": "SciFact Claim Verification (Model B)",
"dataset": "SciFact",
"train_samples": len(train_df),
"valid_samples": len(valid_df),
"epochs": epochs,
"learning_rate": lr,
"batch_size_per_device": per_device_batch_size,
"grad_accum_steps": grad_accum_steps,
"effective_batch_size": per_device_batch_size * grad_accum_steps,
"max_length": max_length,
"class_weights": class_weights.tolist(),
"precision": precision_name,
"gpu_used": gpu_info.get("gpu_name"),
"peak_vram_mb": round(peak_vram_mb, 2),
"training_time_seconds": round(total_training_time, 2),
"best_metrics": best_metrics,
"label2id": LABEL2ID,
"id2label": ID2LABEL,
}
with open(os.path.join(models_dir, "training_config.json"), "w", encoding="utf-8") as f:
json.dump(config, f, indent=2)
# Create README.md inside model dir
readme_content = f"""# TruthLens AI — Claim Verification Model (Model B)
## Overview
- **Model Name:** `{model_name}` (Fine-Tuned)
- **Dataset:** SciFact Scientific Claim Verification (`scifact_train.csv` & `scifact_valid.csv`)
- **Task:** Predict veracity verdict (`SUPPORTS`, `REFUTES`, `NOT_ENOUGH_INFO`) given claim and evidence text.
- **Hardware Acceleration:** {gpu_info.get('gpu_name')} (CUDA {gpu_info.get('cuda_version')})
- **Precision:** Mixed Precision FP16
## Performance Metrics (SciFact Validation Set: {len(valid_df)} samples)
- **Accuracy:** {best_metrics['accuracy']*100:.2f}%
- **Macro F1:** {best_metrics['macro_f1']*100:.2f}%
- **Weighted F1:** {best_metrics['weighted_f1']*100:.2f}%
- **Macro Precision:** {best_metrics['macro_precision']*100:.2f}%
- **Macro Recall:** {best_metrics['macro_recall']*100:.2f}%
### Per-Class Metrics:
- **SUPPORTS:** Precision: {best_metrics['per_class']['SUPPORTS']['precision']*100:.2f}%, Recall: {best_metrics['per_class']['SUPPORTS']['recall']*100:.2f}%, F1: {best_metrics['per_class']['SUPPORTS']['f1']*100:.2f}%
- **REFUTES:** Precision: {best_metrics['per_class']['REFUTES']['precision']*100:.2f}%, Recall: {best_metrics['per_class']['REFUTES']['recall']*100:.2f}%, F1: {best_metrics['per_class']['REFUTES']['f1']*100:.2f}%
- **NOT_ENOUGH_INFO:** Precision: {best_metrics['per_class']['NOT_ENOUGH_INFO']['precision']*100:.2f}%, Recall: {best_metrics['per_class']['NOT_ENOUGH_INFO']['recall']*100:.2f}%, F1: {best_metrics['per_class']['NOT_ENOUGH_INFO']['f1']*100:.2f}%
## Training Configuration
- **Epochs:** {epochs}
- **Learning Rate:** {lr}
- **Per-Device Batch Size:** {per_device_batch_size} (Grad Accum: {grad_accum_steps} -> Effective Batch Size: {per_device_batch_size*grad_accum_steps})
- **Max Length:** {max_length} tokens
- **Peak VRAM:** {peak_vram_mb:.2f} MB
- **Training Time:** {total_training_time:.2f} seconds
"""
with open(os.path.join(models_dir, "README.md"), "w", encoding="utf-8") as f:
f.write(readme_content)
# 9. Plot and Save Confusion Matrix
cm_path = os.path.join(reports_dir, "confusion_matrix.png")
plot_and_save_confusion_matrix(
best_metrics["confusion_matrix"],
class_names=["SUPPORTS", "REFUTES", "NOT_ENOUGH_INFO"],
output_path=cm_path,
title="SciFact Transformer Model B — Confusion Matrix"
)
# 10. Perform Detailed Error Analysis
print("\nCompiling error analysis...")
error_list = []
for i in range(len(valid_df)):
actual = best_targets[i]
pred = best_preds[i]
if actual != pred:
conf = float(best_probs[i][pred])
actual_str = ID2LABEL[actual]
pred_str = ID2LABEL[pred]
error_list.append({
"claim": valid_df.iloc[i]["claim"],
"evidence": valid_df.iloc[i]["evidence"],
"actual_label": actual_str,
"predicted_label": pred_str,
"confidence": round(conf, 4),
})
df_errors = pd.DataFrame(error_list).sort_values("confidence", ascending=False)
errors_csv = os.path.join(reports_dir, "verification_errors.csv")
df_errors.to_csv(errors_csv, index=False)
print(f"Saved {len(df_errors)} error cases to: {errors_csv}")
# 11. Save Hardware Performance Log
hw_report = os.path.join(reports_dir, "training_hardware.md")
log_hardware_metrics(
output_path=hw_report,
gpu_info=gpu_info,
training_time_sec=total_training_time,
total_samples=len(train_df),
epochs=epochs,
peak_vram_mb=peak_vram_mb,
extra_notes=f"- Model Architecture: `{model_name}`\n- Sequence Length: {max_length}\n- Effective Batch Size: {per_device_batch_size*grad_accum_steps}\n- Class Weighting: Applied\n"
)
# 12. Save Validation Results Report
res_report = os.path.join(reports_dir, "scifact_verification_results.md")
cm = best_metrics["confusion_matrix"]
res_md = f"""# TruthLens AI — SciFact Claim Verification Model Results
**Model:** `{model_name}` Fine-Tuned for Evidence Verification (Model B)
**Hardware:** {gpu_info.get('gpu_name')} (CUDA {gpu_info.get('cuda_version')})
**Dataset:** SciFact ({len(train_df)} train pairs, {len(valid_df)} validation pairs)
**Date:** September 2026
**Status:** Completed
---
## 1. Primary Metrics (Validation Set)
| Metric | Score |
| :--- | :--- |
| **Accuracy** | **{best_metrics['accuracy']*100:.2f}%** |
| **Macro F1 Score** | **{best_metrics['macro_f1']*100:.2f}%** |
| **Weighted F1 Score** | **{best_metrics['weighted_f1']*100:.2f}%** |
| **Macro Precision** | **{best_metrics['macro_precision']*100:.2f}%** |
| **Macro Recall** | **{best_metrics['macro_recall']*100:.2f}%** |
---
## 2. Per-Class Breakdown
| Class | Precision | Recall | F1 Score | Validation Count |
| :--- | :--- | :--- | :--- | :--- |
| **`SUPPORTS`** | {best_metrics['per_class']['SUPPORTS']['precision']*100:.2f}% | {best_metrics['per_class']['SUPPORTS']['recall']*100:.2f}% | **{best_metrics['per_class']['SUPPORTS']['f1']*100:.2f}%** | {best_metrics['per_class']['SUPPORTS']['support']} |
| **`REFUTES`** | {best_metrics['per_class']['REFUTES']['precision']*100:.2f}% | {best_metrics['per_class']['REFUTES']['recall']*100:.2f}% | **{best_metrics['per_class']['REFUTES']['f1']*100:.2f}%** | {best_metrics['per_class']['REFUTES']['support']} |
| **`NOT_ENOUGH_INFO`** | {best_metrics['per_class']['NOT_ENOUGH_INFO']['precision']*100:.2f}% | {best_metrics['per_class']['NOT_ENOUGH_INFO']['recall']*100:.2f}% | **{best_metrics['per_class']['NOT_ENOUGH_INFO']['f1']*100:.2f}%** | {best_metrics['per_class']['NOT_ENOUGH_INFO']['support']} |
---
## 3. Confusion Matrix
| Actual \\ Predicted | SUPPORTS | REFUTES | NOT_ENOUGH_INFO | Total |
| :--- | :--- | :--- | :--- | :--- |
| **SUPPORTS** | {cm[0][0]} | {cm[0][1]} | {cm[0][2]} | {best_metrics['per_class']['SUPPORTS']['support']} |
| **REFUTES** | {cm[1][0]} | {cm[1][1]} | {cm[1][2]} | {best_metrics['per_class']['REFUTES']['support']} |
| **NOT_ENOUGH_INFO** | {cm[2][0]} | {cm[2][1]} | {cm[2][2]} | {best_metrics['per_class']['NOT_ENOUGH_INFO']['support']} |
---
## 4. Hardware and Throughput Summary
- **Execution Time:** {total_training_time:.2f} seconds ({total_training_time/60:.2f} minutes)
- **Peak VRAM Allocated:** {peak_vram_mb:.2f} MB (within 4096 MB laptop GPU ceiling)
- **Throughput:** {(len(train_df)*epochs)/total_training_time:.2f} samples/second
"""
with open(res_report, "w", encoding="utf-8") as f:
f.write(res_md)
print(f"Saved results report -> {res_report}")
return {
"best_metrics": best_metrics,
"peak_vram_mb": peak_vram_mb,
"training_time_sec": total_training_time,
"model_dir": models_dir,
}
if __name__ == "__main__":
root_p = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
train_scifact_verification_model(root_p)