add todo for self-gen logits

README.md

@@ -105,8 +105,8 @@ ctx_model_max_len = 2**13
 # load via a custom function because of custom chat_template
 tokenizer = get_tokenizer(base_model_name)
 ctx_tokenizer = get_tokenizer(base_model_name)
-ds = def get_tokenized_dataset(
-    ds_name
+ds = get_tokenized_dataset(
+    ds_name,
     split="train",
     base_model_max_len=base_model_max_len,
     tokenizer=tokenizer,

data/self_generate_qa.py

@@ -1,6 +1,5 @@
 import argparse
 import os
-import random
 from glob import glob
 
 import pandas as pd
@@ -160,11 +159,12 @@ def self_generate(
     messages = create_messages(ctxs, questions, args.vllm_model, SYSTEM_TEMPLATE)
 
     print(f"Generating from {len(messages)} contexts")
+    # TODO (distillation): make vllm outputs logits here too
     completions = llm.chat(
         messages,
         sampling_params=SamplingParams(
             max_tokens=2048,
-            temperature=1.0,
+            temperature=1.0,  # TODO: lower the temp
         ),
     )
 

src/ctx_to_lora/data/processing.py

@@ -592,6 +592,7 @@ def get_tokenized_dataset(
     set_format: str | None = None,
     streaming: bool = False,
 ) -> dict[str, Any]:
+    # TODO (distillation, Tan): make this works with pre-computed logits
     assert not use_kl_loss, "KL loss is deprecated"
     if add_repeat_prompt:
         assert repeat_prob > 0, f"add_repeat_prompt is set but repeat_prob = 0"

src/ctx_to_lora/trainer.py

@@ -10,7 +10,7 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
 logger = logging.getLogger()
 
 
-# TODO (distillation): implement
+# TODO (distillation, Shin): implement
 class DistillationTrainer(Trainer): ...