File size: 4,362 Bytes
fe775e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""Continue training from saved model with augmented data."""
import json
import os
import sys
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
from datasets import Dataset
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import RAW_DIR, MODELS_DIR, MAX_SEQ_LEN, LABEL_MAP, BATCH_SIZE, LEARNING_RATE
FINAL_MODEL_DIR = os.path.join(MODELS_DIR, "final_model")
def load_all_jsonl():
examples = []
for fname in os.listdir(RAW_DIR):
if not fname.endswith(".jsonl") or "extra" in fname:
continue
fpath = os.path.join(RAW_DIR, fname)
with open(fpath) as f:
for line in f:
line = line.strip()
if not line: continue
obj = json.loads(line)
text = obj.get("text", "").strip()
label_str = obj.get("label", "")
if text and label_str in LABEL_MAP:
examples.append({"text": text, "label": LABEL_MAP[label_str]})
return examples
def main():
print("Loading saved model...")
tokenizer = AutoTokenizer.from_pretrained(FINAL_MODEL_DIR)
model = AutoModelForSequenceClassification.from_pretrained(FINAL_MODEL_DIR)
print("Loading augmented dataset...")
examples = load_all_jsonl()
print(f"Total examples: {len(examples)}")
# Check label balance
labels = [ex["label"] for ex in examples]
print(f" GENERIC: {labels.count(0)}")
print(f" SEMANTIC: {labels.count(1)}")
dataset = Dataset.from_list(examples)
splits = dataset.train_test_split(test_size=0.15, seed=42)
def tokenize(examples):
return tokenizer(examples["text"], padding="max_length",
truncation=True, max_length=MAX_SEQ_LEN)
tokenized = splits.map(tokenize, batched=True)
tokenized = tokenized.remove_columns(["text"])
tokenized = tokenized.rename_column("label", "labels")
tokenized.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
print("\nContinuing training for 1 epoch...")
training_args = TrainingArguments(
output_dir=os.path.join(MODELS_DIR, "checkpoints_v2"),
eval_strategy="steps",
eval_steps=100,
save_strategy="steps",
save_steps=200,
logging_steps=25,
learning_rate=5e-6, # Lower LR for continued training
per_device_train_batch_size=BATCH_SIZE,
per_device_eval_batch_size=BATCH_SIZE * 2,
num_train_epochs=1,
weight_decay=0.01,
warmup_ratio=0.05,
fp16=torch.cuda.is_available(),
save_total_limit=1,
load_best_model_at_end=True,
metric_for_best_model="accuracy",
greater_is_better=True,
report_to="none",
seed=42,
dataloader_num_workers=2,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized["train"],
eval_dataset=tokenized["test"],
compute_metrics=lambda p: (
{"accuracy": (p.predictions.argmax(-1) == p.label_ids).mean()}
),
)
trainer.train()
# Evaluate on specific problem cases
print("\nEvaluating problem cases:")
model.eval()
problem_queries = [
"I love spicy food",
"my name is John",
"मेरा नाम रवि है",
"नमस्ते",
"hello",
"मुझे कॉफी पसंद है",
"I work as a software engineer",
"my favorite color is blue",
]
for query in problem_queries:
inputs = tokenizer(query, return_tensors="pt", padding="max_length",
truncation=True, max_length=MAX_SEQ_LEN)
with torch.no_grad():
logits = model(**inputs).logits
probs = torch.nn.functional.softmax(logits, dim=-1).numpy()[0]
pred = "SEMANTIC" if probs[1] > probs[0] else "GENERIC"
print(f" [{pred:8s}] gen={probs[0]:.3f} sem={probs[1]:.3f} \"{query}\"")
# Save model again
trainer.save_model(FINAL_MODEL_DIR)
tokenizer.save_pretrained(FINAL_MODEL_DIR)
print(f"\nModel saved to {FINAL_MODEL_DIR}")
if __name__ == "__main__":
main()
|