Low Level Design

Low-Level Design (LLD)

This design is derived from DistilBERT_Colab_Fixed_2026.py. The diagram groups module-level functions as ProjectFunctions; it does not claim the source defines that Python class. Library types are existing Hugging Face abstractions.

Module dependencies

classDiagram
    class ProjectFunctions {
      <<module-level functions>>
      +tokenize_function(examples)
      +compute_metrics(eval_pred)
      +create_training_arguments(kwargs)
      +create_trainer(kwargs)
      +classify_query(query)
    }
    class DatasetDict {
      +train
      +validation
    }
    class AutoTokenizer {
      +from_pretrained(model_name)
      +__call__(texts)
    }
    class DataCollatorWithPadding {
      +__call__(features)
    }
    class TrainingArguments {
      +num_train_epochs
      +learning_rate
      +eval_strategy
    }
    class Trainer {
      +train()
      +evaluate()
      +predict(dataset)
      +save_model(path)
    }
    class AutoModelForSequenceClassification {
      +forward(input_ids, attention_mask)
    }
    ProjectFunctions --> DatasetDict : prepares and maps
    ProjectFunctions --> AutoTokenizer : tokenizes
    ProjectFunctions --> TrainingArguments : creates
    ProjectFunctions --> Trainer : configures and invokes
    Trainer --> AutoModelForSequenceClassification : trains
    Trainer --> DataCollatorWithPadding : batches

Loading

Control flow

flowchart TD
    Start[Import and seed] --> Dataset[RAW_DATA + LABEL2ID]
    Dataset --> Split[train_test_split stratified]
    Split --> Map[dataset.map tokenize_function]
    Map --> Args[create_training_arguments]
    Args --> Build[create_trainer]
    Build --> Run[trainer.train]
    Run --> Save[trainer.save_model + tokenizer.save_pretrained]
    Save --> Eval[trainer.evaluate + trainer.predict]
    Eval --> Classify[classify_query]
    Classify --> Freeze[Freeze model_frozen.distilbert parameters]
    Freeze --> Train2[frozen_trainer.train + evaluate]
    Train2 --> Compare[Print metrics comparison]

Loading

Key functions

Function	Inputs	Output / side effect
tokenize_function(examples)	Dataset batch with text	Encoded input_ids and masks; truncates to 128
compute_metrics(eval_pred)	Logits and true class IDs	Accuracy and weighted F1
create_training_arguments(**kwargs)	Train config	TrainingArguments with version-sensitive warmup and evaluation key selection
create_trainer(**kwargs)	Model, datasets, args, collator, metric hook	Trainer with tokenizer/processing_class chosen by introspection
classify_query(query)	A raw customer-support message	Category, top probability, four class scores

State and data types

• RAW_DATA: list of 30 (text, label_name) tuples.
• LABEL_NAMES, LABEL2ID, ID2LABEL: fixed mapping for four categories.
• DatasetDict: train and validation splits.
• tokenized_dataset: tensors for input_ids, attention_mask, and label.
• model, trainer: full fine-tuning objects; model_frozen, frozen_trainer used for comparison.

The separate Phi-2 source is too incomplete to establish a trustworthy executable call graph; see Phi2-SFT-DPO-Design for its intended design.