| """ |
| train_context.py |
| |
| Train the HV Context model on a conversation dataset. |
| |
| If you have your own turns (list of (topic, user, assistant)), |
| edit the load_turns() function to read from your source. |
| |
| Run: |
| python train_context.py |
| python train_context.py --output ./my_context_model |
| """ |
|
|
| import argparse |
| import time |
| from pathlib import Path |
|
|
| from hv_context import ( |
| HVContext, generate_conversation, D, SEED, |
| ) |
|
|
| OUTPUT_DIR = "context_model" |
|
|
|
|
| def load_turns(): |
| """Replace this with your own data loader. |
| |
| Returns a list of (topic, user_text, assistant_text) tuples. |
| The topic label is only used for evaluation -- it is not |
| consumed by the fit() method. |
| """ |
| return generate_conversation(n_per_topic=10, seed=SEED) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Train HV Context.") |
| parser.add_argument("--output", type=str, default=OUTPUT_DIR) |
| parser.add_argument("--dim", type=int, default=D) |
| args = parser.parse_args() |
|
|
| print(f"Loading turns...") |
| turns = load_turns() |
| print(f" {len(turns)} turns, {len(set(t for t, _, _ in turns))} topics") |
|
|
| print(f"Training...") |
| t0 = time.perf_counter() |
| model = HVContext(D=args.dim, seed=SEED) |
| model.fit(turns) |
| dt = time.perf_counter() - t0 |
|
|
| print(f" vocabulary: {len(model.vocab_index)} words") |
| print(f" codebook: {model.K.size / 8 / 1024:.1f} KB") |
| print(f" bank: {model.bank_int8.nbytes / 1024:.1f} KB") |
| print(f" time: {dt*1000:.0f} ms") |
|
|
| print(f"Saving to {args.output}/") |
| model.save(args.output) |
| print("Done.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
