{"id":28470785,"url":"https://github.com/kvignesh1420/cot-icl-lab","last_synced_at":"2025-07-01T20:32:51.192Z","repository":{"id":292778380,"uuid":"938998653","full_name":"kvignesh1420/cot-icl-lab","owner":"kvignesh1420","description":"[ACL 2025] CoT-ICL Lab: A Synthetic Framework for Studying Chain-of-Thought Learning from In-Context Demonstrations","archived":false,"fork":false,"pushed_at":"2025-05-23T02:11:15.000Z","size":568,"stargazers_count":10,"open_issues_count":0,"forks_count":1,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-06-07T10:06:21.882Z","etag":null,"topics":["chain-of-thought","graphs","in-context-learning","llm-inference","llm-training","transformers"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2502.15132","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/kvignesh1420.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2025-02-25T20:38:24.000Z","updated_at":"2025-05-31T13:34:23.000Z","dependencies_parsed_at":"2025-05-12T05:29:59.249Z","dependency_job_id":"da4b128b-d376-4cf4-8c9b-6afff84fec90","html_url":"https://github.com/kvignesh1420/cot-icl-lab","commit_stats":null,"previous_names":["kvignesh1420/cot-icl-lab"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/kvignesh1420/cot-icl-lab","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/kvignesh1420%2Fcot-icl-lab","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/kvignesh1420%2Fcot-icl-lab/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/kvignesh1420%2Fcot-icl-lab/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/kvignesh1420%2Fcot-icl-lab/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/kvignesh1420","download_url":"https://codeload.github.com/kvignesh1420/cot-icl-lab/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/kvignesh1420%2Fcot-icl-lab/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":263033197,"owners_count":23403114,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["chain-of-thought","graphs","in-context-learning","llm-inference","llm-training","transformers"],"created_at":"2025-06-07T10:06:25.459Z","updated_at":"2025-07-01T20:32:51.179Z","avatar_url":"https://github.com/kvignesh1420.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n  \u003ch1\u003eCoT-ICL Lab\u003c/h1\u003e\n  \u003cp\u003e A Synthetic Framework for Studying Chain-of-Thought Learning from In-Context Demonstrations\n \u003c/p\u003e\n\u003c/div\u003e\n\u003cbr\u003e\n\n![](assets/cot_icl_intro.png)\n\n****************************************************************\n\n## Setup\n\n- Create a virtual environment and install the package.\n\n```bash\n$ python3.12 -m venv .venv\n$ source .venv/bin/activate\n(.venv) $ pip install -e .\n```\n\n- Run unit tests as a sanity check.\n\n```bash\n(.venv) $ pytest\n```\n\n- (Development) Run ruff + isort fixes to sanitize the code changes.\n\n```bash\n(.venv) $ ./beautify.sh\n```\n\n## Getting Started\n\nOur framework serves as a test bed to generate synthetic tokenized datasets for training and evaluating transformer models. We do so by using `DAG` and `TokenProcessor` classes. These can be configured directly by the `Args` dataclass. For example:\n\n```py\n\nfrom tokenized_cot_icl.core.args import Args\nfrom tokenized_cot_icl.core.data import TokenizedDataset\n\nargs = Args(\n      vocab_size=1024,\n      n_inputs=4,\n      n_parents=2,\n      chain_length=3,\n      n_examples=1,\n      enable_cot=True,\n      prompt_strategy=\"cot\",\n      activation=\"leaky_relu\",\n      n_tasks=10,\n)\n\ndataset = TokenizedDataset(args=args)\nprint(dataset[0])\n```\n\nThe above item in the dataset is as follows:\n\n```py\n{\n    'adj_list': tensor([[0, 2], [4, 3], [5, 3]]),\n    'attention_mask': tensor([1, 1, 1, 1, 1, 1, 1]),\n    'input_ids': tensor([ 556,  197, 1002,  867,  240,  466,  217]),\n    'labels': tensor([-100, -100, -100, -100,  240,  466,  217]),\n    'cot_eval':\n          {\n                'attention_mask': tensor([1, 1, 1, 1]),\n                'input_ids': tensor([ 556,  197, 1002,  867]),\n                'last_example_cot': tensor([240, 466, 217])\n          }\n}\n```\n\n### Understanding the DAG structure\n\nLet's break down the result above to understand the DAG structure. Consider $4$ input tokens ($x_1, x_2, x_3, x_4$) and $3$ chain tokens ($y_1, y_2, y_3$) for the single example above.\n\nThe `'adj_list': tensor([[0, 2], [4, 3], [5, 3]])` (based on zero-indexing) indicates that the parent tokens for the chain tokens are as follows:\n\n\u003cdiv align=\"center\"\u003e\n\n| Chain Token | Parent Tokens    |\n|-------------|------------------|\n| $y_1$  | $\\{x_1, x_3\\}$ |\n| $y_2$ | $\\{y_1, x_4\\}$ |\n| $y_3$ | $\\{y_2, x_4\\}$ |\n\n\u003c/div\u003e\n\n\n\u003e[!NOTE]\n\u003e The TokenCoverage metric introduced in the paper relies on the uniqueness of chain tokens in the entire dataset and depends heavily on the \"vocab_size\" and \"activation\". Thus controlling the difficulty of the tasks.\n\n## Models\n\nWe leverage the HuggingFace [transformers](https://github.com/huggingface/transformers) library to create custom Llama models and expose a `MODEL_REGISTRY` to register new model families.\n\n```py\n# src/tokenized_cot_icl/core/models.py\n\nMODEL_REGISTRY = {\"llama\": create_llama_model}\n```\n\n\u003e[!TIP]\n\u003e Users can register the creation function for models of their choice from the `transformers` library to explore new architectures and validate ideas.\n\n\n## Training\n\n### Setting the `TASK_CARD`\n\nTo make it suitable for bulk launching the experiments, we rely on a `TASK_CARD` to collate all the args. For instance, to train a model with the args as per the above example, we do:\n\n```py\n# src/tokenized_cot_icl/core/task_card.py\n\ndef custom_task_card() -\u003e Dict[int, Args]:\n    \"\"\"A custom task card.\"\"\"\n      args = Args(...) # set as needed\n    return {0: args}\n\n# set the dictionary\nTASK_CARD = custom_task_card()\n```\n\n### Launch the DDP Training\n\nThe `TASK_CARD` allows us to index into the experimental config of our choice and launch the torch distributed data parallel (DDP) training runs. For example:\n\n```bash\n(.venv) $ cd src\n(.venv) $ export NUM_NODES=1 # change as needed\n(.venv) $ export LOCAL_WORLD_SIZE=4 # change as needed\n(.venv) $ torchrun --nnodes=$NUM_NODES --nproc-per-node=$LOCAL_WORLD_SIZE -m tokenized_cot_icl.core.train --task_card_key 0\n```\n\n### Metric Logging\n\n- By default, we use `metric_logger=\"stdout\"` in `Args` and log the metrics/params to `STDOUT`.\n- We also support logging to an [MLFlow](https://mlflow.org/docs/latest/tracking.html) tracking server by setting the `MLFLOW_SERVICE_URL` environment variable and using `Args(metric_logger=\"mlflow\")`.\n\n### Liger-Kernels\n\nUsers can also apply the [Liger-Kernel](https://github.com/linkedin/Liger-Kernel) optimizations to patch the llama models by setting `Args(use_liger_kernels=True)` and speed up the training runs.\n\n```bash\n(.venv) $ pip install liger-kernel # install suitable version\n```\n\n## Inference with vLLM/SGLang\n\nIn addition to using the `transformers.GenerationConfig` for small scale inference during the training runs, we also support [vLLM](https://github.com/vllm-project/vllm) and [SGLang](https://github.com/sgl-project/sglang) based evaluation of the trained model (or model checkpoints) to analyze the predictions.\n\n```bash\n(.venv) $ pip install vllm # install suitable version\n(.venv) $ pip install sglang # install suitable version\n```\n\nWe provide an easy to extend example for calculating the answer token prediction accuracy as follows:\n\n```bash\n# for vllm\n(.venv) $ cd src \u0026\u0026 python tokenized_cot_icl/inference/vllm/evaluator.py \\\n                        --model_base_dir /opt/cot-icl-lab/run_name \\\n                        --checkpoint final  # either final or 1000, 2000 etc.\n\n# for sglang\n(.venv) $ cd src \u0026\u0026 python tokenized_cot_icl/inference/sglang/evaluator.py \\\n                        --model_base_dir /opt/cot-icl-lab/run_name \\\n                        --checkpoint final  # either final or 1000, 2000 etc.\n```\n\n\n## License\n\n[MIT License](LICENSE)\n\n## Citation\n\n```bibtex\n@misc{kothapalli2025coticllabsyntheticframework,\n      title={CoT-ICL Lab: A Synthetic Framework for Studying Chain-of-Thought Learning from In-Context Demonstrations}, \n      author={Vignesh Kothapalli and Hamed Firooz and Maziar Sanjabi},\n      year={2025},\n      eprint={2502.15132},\n      archivePrefix={arXiv},\n      primaryClass={cs.CL},\n      url={https://arxiv.org/abs/2502.15132}, \n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fkvignesh1420%2Fcot-icl-lab","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fkvignesh1420%2Fcot-icl-lab","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fkvignesh1420%2Fcot-icl-lab/lists"}