Sync from GitHub via hub-sync

CODE_OF_CONDUCT.md

@@ -0,0 +1,136 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at:
+**akram.alsubari87@gmail.com** or **akram.alsubari@outlook.com**
+
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior, harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.
+
+---
+
+## Contact
+
+For any Code of Conduct concerns or questions, please contact me through emails and linked:
+**www.linkedin.com/in/akram-alsubari**

CONTRIBUTING.md

@@ -0,0 +1,22 @@
+# Contributing to dytr
+
+## How to Contribute
+
+1. Fork the repository
+2. Create a feature branch (`git checkout -b feature/amazing-feature`)
+3. Commit your changes (`git commit -m 'Add amazing feature'`)
+4. Push to branch (`git push origin feature/amazing-feature`)
+5. Open a Pull Request
+
+## Pull Request Process
+
+1. Ensure your code passes all tests
+2. Update documentation if needed
+3. PR will be merged after approval
+
+## Development Setup
+
+```bash
+git clone https://github.com/AAlsubari/dytr.git
+cd dytr
+pip install -e .[dev]

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

MANIFEST.in

@@ -0,0 +1,7 @@
+include LICENSE
+include NOTICE
+include README.md
+include pyproject.toml
+recursive-include src/dytr *.py
+recursive-include src/dytr *.json
+recursive-include src/dytr *.txt

NOTICE

@@ -0,0 +1,15 @@
+dytr - Dynamic Transformer Library
+Copyright 2025 Akram Alsubari
+
+This product includes software developed by:
+Akram Alsubari (http://www.linkedin.com/in/akram-alsubari)
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

README.md

@@ -1,48 +1,15 @@
----
-license: apache-2.0
-language:
-- ar
-- en
-- af
-- fr
-- fa
-- ff
-- fi
-- fj
-- fo
-- fy
-- qu
-- zh
-- za
-- zu
-- he
-- es
-- el
-- ro
-- ta
-- pa
-- hi
-- de
-metrics:
-- perplexity
-- f1
-- accuracy
-base_model:
-- google-bert/bert-base-uncased
-datasets:
-- alsubari/arabic-grammar-errors
----
 # dytr - Dynamic Transformer Library
+dytr is a flexible PyTorch library for multi-task learning with dynamic transformer architectures. Train multiple tasks sequentially or simultaneously while preserving performance on previous tasks through built-in continual learning techniques. it also supports to finetune and modify pretrained model such as bert.
+
+
+
 
 [![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 [![PyPI version](https://badge.fury.io/py/dytr.svg)](https://badge.fury.io/py/dytr)
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/AAlsubari/dytr/blob/main/dytr_bert_finetune_demo.ipynb)
 
-**Build dynamic transformers that learn multiple tasks with supports to load and modify pretrained model such as bert.**
-
-dytr is a flexible PyTorch library for multi-task learning with transformer architectures. Train multiple tasks sequentially or simultaneously while preserving performance on previous tasks through built-in continual learning techniques.
-
+**Build dynamic transformers that learn multiple tasks.**
 ## Why dytr?
 
 - 🎯 **Multi-Task Ready** - Train classification, generation, and sequence tasks in one model
@@ -50,14 +17,12 @@ dytr is a flexible PyTorch library for multi-task learning with transformer arch
 - 🔧 **No Black Box** - Full control over architecture, understand every component
 - ⚡ **Lightweight** - Pure PyTorch, minimal dependencies
 - 📦 **Pretrained Support** - Load BERT, RoBERTa, and more as your encoder backbone and fine tune it on multiple tasks.
-
+- 
 
 ## Architecture Overview
 
-![unnamed](https://cdn-uploads.huggingface.co/production/uploads/645a5b8e333fb1835783d500/TdQZhE7EPeGU9zgYJMKbl.png)
-
-
 
+![unnamed](https://cdn-uploads.huggingface.co/production/uploads/645a5b8e333fb1835783d500/TdQZhE7EPeGU9zgYJMKbl.png)
 
 ## Installation
 
@@ -162,12 +127,12 @@ Load powerful encoders as your backbone and extend them with tasks:
 
 ```python
 from dytr import PretrainedModelLoader
-
+model_name='prajjwal1/bert-tiny'
 loader = PretrainedModelLoader()
-config = ModelConfig(tokenizer_name='bert-base-uncased',per_device_train_batch_size=32,num_train_epochs=3,per_device_eval_batch_size=8,special_tokens={},use_task_adapters=False,use_ewc=True,use_replay=True,use_rotary_embedding=False, training_from_scratch=False)
+config = ModelConfig(tokenizer_name=model_name, per_device_train_batch_size=32, num_train_epochs=3, per_device_eval_batch_size=8, special_tokens={}, use_task_adapters=False, use_ewc=True, use_replay=True, use_rotary_embedding=False, training_from_scratch=False)
 
 # Load pretrained BERT as your encoder
-model = loader.load_pretrained('bert-base-uncased', config)
+model = loader.load_pretrained(model_name, config)
 
 # Now add your own tasks - the model is fully dytr compatible
 class_train = pd.DataFrame(
@@ -229,7 +194,6 @@ train_datasets = {
         classification_task.task_name: (class_dataset, classification_task.training_strategy),
         lm_task.task_name: (lm_dataset, lm_task.training_strategy),
     }
-
 val_datasets = {
         #classification_task.task_name: (class_val_dataset, classification_task.training_strategy)
     }
@@ -278,6 +242,8 @@ config = ModelConfig(
 
 
 
+
+
 ## Who Should Use dytr?
 
 | Audience | Why It Matters |
@@ -315,15 +281,10 @@ config = ModelConfig(
 
 Apache License 2.0
 
-## project links
-- [PyPI Page](https://pypi.org/project/dytr/)
-- [GitHub Repository](https://github.com/AAlsubari/dytr)
-
 ## Author
 
 **Dr. Akram Alsubari**
-- akram.alsubari@outlook.com
-- akram.alsubari87@gmail.com
+
 
 ## Contributing
 
@@ -338,4 +299,4 @@ For questions, issues, or suggestions:
 - 🎓 **Research Interests**: Natural Language Processing, Deep Learning, Transformers, Continual Learning, Multi-Task Learning, Large Language Models
 ---
 
-**Build once. Learn multiple tasks. Never forget.**
+**Build once. Learn multiple tasks. Never forget.**

SECURITY.md

@@ -0,0 +1,91 @@
+# Security Policy
+
+## Supported Versions
+
+Currently supported versions with security updates:
+
+| Version | Supported |
+|---------|-----------|
+| 0.1.x   | ✅ Yes    |
+| < 0.1   | ❌ No     |
+
+## Reporting a Vulnerability
+
+### How to Report
+
+**DO NOT** create public GitHub issues for security vulnerabilities.
+
+Please report security vulnerabilities by email to:
+
+📧 **akram.alsubari@outlook.com**
+
+📧 **akram.alsubari87@gmail.com** (alternative)
+
+### What to Include
+
+When reporting a vulnerability, please include:
+
+- Description of the vulnerability
+- Steps to reproduce (code examples if possible)
+- Affected versions
+- Potential impact
+- Any suggested fixes (optional)
+
+### What to Expect
+
+| Step | Timeline | Description |
+|------|----------|-------------|
+| 1 | Within 48 hours | Acknowledgment of receipt |
+| 2 | Within 5-7 days | Initial assessment and confirmation |
+| 3 | Within 14 days | Fix development or mitigation plan |
+| 4 | Within 30 days | Release of security update |
+
+### Vulnerability Acceptance
+
+If the vulnerability is accepted:
+- You will receive credit in the release notes (unless you prefer to remain anonymous)
+- A security advisory will be published
+- A patch will be released for supported versions
+
+### Vulnerability Decline
+
+If the vulnerability is declined:
+- You will receive a detailed explanation why
+- The issue may be moved to a public discussion if appropriate
+- No further action will be taken
+
+## Security Best Practices for Users
+
+When using dytr:
+
+1. **Keep updated**: Always use the latest version
+2. **Validate inputs**: Sanitize any untrusted input before passing to the model
+3. **Use virtual environments**: Isolate dytr installations
+4. **Review code changes**: Check pull requests before merging
+
+## Security Updates
+
+Security updates will be released as:
+
+- **Patch versions** (0.1.x → 0.1.y) for critical fixes
+- **Minor versions** (0.x → 0.y) for non-critical fixes
+
+Watch the repository on GitHub to receive security update notifications.
+
+## Responsible Disclosure
+
+We follow responsible disclosure practices:
+- Reporters will be acknowledged (unless anonymous)
+- Fixes will be released before public disclosure
+- Coordinated disclosure with reporters
+
+## Contact
+
+**Security Contact:** akram.alsubari@outlook.com
+
+
+**Alternative Contact:** akram.alsubari87@gmail.com
+
+## Acknowledgment
+
+We thank the security research community for helping keep dytr secure.

config.json

@@ -1,20 +0,0 @@
-{
-  "embed_dim": 256,
-  "num_layers": 4,
-  "num_heads": 4,
-  "head_dim": 64,
-  "ff_mult": 4,
-  "max_seq_len": 512,
-  "dropout": 0.1,
-  "vocab_size": 32000,
-  "tokenizer_name": "asafaya/bert-mini-arabic",
-  "use_rotary_embedding": false,
-  "use_task_adapters": false,
-  "adapter_bottleneck": 64,
-  "tasks": [
-    "sentiment",
-    "error_detection",
-    "ner_detection",
-    "generation"
-  ]
-}

finetune_bert_mini_arabic_mltitasks_and_generation.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:814785ecd561cc5d5c82bba2a63a5ed0e8ef571af9e012a4a06667e7c7a453f0
-size 100112722

pyproject.toml

@@ -0,0 +1,95 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel>=0.37.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "dytr"
+version = "0.1.0"
+description = "dytr: Dynamic Transformer for Multi-Task Learning with Continual Learning Support"
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "Apache-2.0"}
+authors = [
+    {name = "Akram Alsubari", email = "akram.alsubari@outlook.com"}
+]
+maintainers = [
+    {name = "Akram Alsubari", email = "akram.alsubari87@gmail.com"}
+]
+keywords = [
+    "transformer",
+    "multi-task-learning",
+    "continual-learning",
+    "pytorch",
+    "nlp",
+    "deep-learning",
+    "attention",
+    "neural-networks",
+    "machine-learning",
+    "dytr"
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Information Technology",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Information Analysis",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: Text Processing :: Linguistic",
+    "Operating System :: OS Independent"
+]
+
+dependencies = [
+    "torch>=1.10.0",
+    "numpy>=1.19.0",
+    "pandas>=1.3.0",
+    "scikit-learn>=0.24.0",
+    "tqdm>=4.62.0",
+    "requests>=2.25.0"
+]
+
+[project.optional-dependencies]
+transformers = [
+    "transformers>=4.0.0"
+]
+dev = [
+    "pytest>=6.0",
+    "pytest-cov>=2.0",
+    "black>=21.0",
+    "flake8>=3.9",
+    "mypy>=0.910",
+    "build>=0.7",
+    "twine>=3.4",
+    "isort>=5.9.0"
+]
+docs = ["sphinx>=4.0", "sphinx-rtd-theme>=0.5"]
+all = [
+    "transformers>=4.0.0",
+    "pytest>=6.0",
+    "sphinx>=4.0"
+]
+
+[project.urls]
+"Bug Reports" = "https://github.com/AAlsubari/dytr"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+packages = {find = {where = ["src"], include = ["dytr*"]}}
+include-package-data = true
+
+[tool.setuptools.package-data]
+dytr = [
+    "py.typed",
+    "**/*.json",
+    "**/*.txt",
+]
+
+[tool.setuptools.dynamic]
+version = {attr = "dytr.__version__"}

requirements.txt

@@ -0,0 +1,11 @@
+# Core dependencies for Dynamic Transformers
+
+torch>=1.10.0
+numpy>=1.19.0
+pandas>=1.3.0
+scikit-learn>=0.24.0
+tqdm>=4.62.0
+requests>=2.25.0
+
+# Optional dependencies (commented out)
+# transformers>=4.0.0

setup.py

@@ -0,0 +1,116 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Setup script for Dynamic Transformers library.
+
+This file provides backward compatibility with older pip versions
+and can be used as an alternative to pyproject.toml.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from setuptools import setup, find_packages
+import os
+import sys
+
+# Read the contents of README file
+this_directory = os.path.abspath(os.path.dirname(__file__))
+with open(os.path.join(this_directory, 'README.md'), encoding='utf-8') as f:
+    long_description = f.read()
+
+# Read version from __init__.py
+version = {}
+with open(os.path.join(this_directory, 'src', 'dytr', '__init__.py'), encoding='utf-8') as f:
+    for line in f:
+        if line.startswith('__version__'):
+            exec(line, version)
+            break
+
+setup(
+    name="dytr",
+    version=version.get('__version__', '0.1.0'),
+    author="Akram Alsubari",
+    author_email="akram.alsubari@outlook.com",
+    maintainer="Akram Alsubari",
+    maintainer_email="akram.alsubari87@gmail.com",
+    description="Dynamic Transformer for Multi-Task Learning with Continual Learning Support",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="",  # No GitHub URL
+    project_urls={
+        "Bug Reports": "http://www.linkedin.com/in/akram-alsubari",
+        "Source": "",
+        "Documentation": "",
+    },
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Education",
+        "Intended Audience :: Information Technology",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Scientific/Engineering :: Information Analysis",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+        "Topic :: Text Processing :: Linguistic",
+        "Operating System :: OS Independent",
+    ],
+    python_requires=">=3.8",
+    install_requires=[
+        "torch>=1.10.0",
+        "numpy>=1.19.0",
+        "pandas>=1.3.0",
+        "scikit-learn>=0.24.0",
+        "tqdm>=4.62.0",
+        "requests>=2.25.0",
+    ],
+    extras_require={
+        "transformers": ["transformers>=4.0.0"],
+        "dev": [
+            "pytest>=6.0",
+            "pytest-cov>=2.0",
+            "black>=21.0",
+            "flake8>=3.9",
+            "mypy>=0.910",
+            "build>=0.7",
+            "twine>=3.4",
+        ],
+        "docs": [
+            "sphinx>=4.0",
+            "sphinx-rtd-theme>=0.5",
+        ],
+        "all": [
+            "transformers>=4.0.0",
+            "pytest>=6.0",
+            "sphinx>=4.0",
+        ],
+    },
+    include_package_data=True,
+    zip_safe=False,
+    entry_points={
+        "console_scripts": [
+            "dytr-train=dytr.cli.train:main",
+            "dytr-export=dytr.cli.export:main",
+        ],
+    },
+)

src/dytr/__init__.py

@@ -0,0 +1,102 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+dytr - Dynamic Transformer for Multi-Task Learning with Continual Learning Support
+
+A flexible PyTorch library for multi-task learning with transformer architectures.
+
+
+"""
+
+__version__ = "0.1.0"
+__author__ = "Akram Alsubari"
+__email__ = "akram.alsubari@outlook.com"
+__license__ = "Apache 2.0"
+
+from dytr.core.config import ModelConfig, TaskConfig, TrainingStrategy
+from dytr.core.exporter import ModelExporter
+
+# Core model components
+from dytr.core.model import DynamicTransformer
+
+# Memory/Continual learning components
+from dytr.memory.ewc import EWC
+from dytr.memory.replay import ReplayBuffer
+from dytr.tokenization.download_manager import (
+    DownloadManager,
+    download_file,
+    download_tokenizer_files,
+    get_url_from_HF,
+    process_vocab_text,
+)
+
+# Tokenizer
+from dytr.tokenization.simple_tokenizer import SimpleTokenizer
+# Pretrained model loading
+#from dytr.pretrained import PretrainedModelLoader
+from dytr.pretrained.loader import PretrainedModelLoader
+from dytr.training.data import MultiTaskDataset, TaskAwareBatchSampler, collate_fn
+from dytr.training.dataset import SingleDatasetProcessing
+
+# Training components
+from dytr.training.trainer import Trainer
+
+# Utility functions
+from dytr.training.utils import process_qa_dataset, set_seed
+from dytr.utils.caching import CacheManager, get_cache_manager
+from dytr.utils.logging import get_logger, set_log_level, disable_logging
+
+
+# Define what gets imported with "from dynamic_transformers import *"
+__all__ = [
+    # Core classes
+    "DynamicTransformer",
+    "ModelConfig",
+    "TaskConfig",
+    "TrainingStrategy",
+    "Trainer",
+    "ModelExporter",
+    # Dataset classes
+    "SingleDatasetProcessing",
+    "MultiTaskDataset",
+    "TaskAwareBatchSampler",
+    "collate_fn",
+    # Continual learning
+    "EWC",
+    "ReplayBuffer",
+    # Pretrained Models
+    "PretrainedModelLoader",
+    # Tokenizer
+    "SimpleTokenizer",
+    "DownloadManager",
+    "download_file",
+    "download_tokenizer_files",
+    "get_url_from_HF",
+    "process_vocab_text",
+    # Utilities
+    "set_seed",
+    "process_qa_dataset",
+    "CacheManager",
+    "get_cache_manager",
+    "get_logger",
+    "set_log_level", 
+    "disable_logging",
+    # Version info
+    "__version__",
+    "__author__",
+    "__email__",
+    "__license__",
+]
+

src/dytr/cli/__init__.py

@@ -0,0 +1,28 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Command-line interface for Dynamic Transformers.
+
+This module provides CLI tools for training and exporting models.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from dytr.cli.export import main as export_main
+from dytr.cli.train import main as train_main
+
+__all__ = ["train_main", "export_main"]
+

src/dytr/cli/export.py

@@ -0,0 +1,90 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Export CLI for Dynamic Transformers.
+
+This module provides a command-line interface for exporting single-task models.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Export a single-task model from a multi-task model"
+    )
+
+    parser.add_argument(
+        "--model_path", type=str, required=True, help="Path to the saved multi-task model"
+    )
+
+    parser.add_argument("--task_name", type=str, required=True, help="Name of the task to export")
+
+    parser.add_argument(
+        "--output_path", type=str, required=True, help="Path to save the exported model"
+    )
+
+    parser.add_argument(
+        "--device",
+        type=str,
+        choices=["cuda", "cpu"],
+        default=None,
+        help="Device to use for loading the model",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    """Main export function."""
+    args = parse_args()
+
+    print("=" * 60)
+    print("Dynamic Transformers - Export CLI")
+    print("=" * 60)
+
+    print(f"\nLoading model from: {args.model_path}")
+    print(f"Exporting task: {args.task_name}")
+    print(f"Saving to: {args.output_path}")
+
+    # Import library components
+    from dytr import DynamicTransformer
+
+    # Load model
+    model = DynamicTransformer.load_model(args.model_path, device=args.device)
+
+    if model is None:
+        print("Error: Failed to load model")
+        sys.exit(1)
+
+    # Export task
+    exporter = model.get_exporter()
+    exported_model = exporter.export_single_task(args.task_name, args.output_path)
+
+    print(f"\nSuccessfully exported {args.task_name} to {args.output_path}")
+
+    print("\n" + "=" * 60)
+    print("Export completed successfully!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

src/dytr/cli/train.py

@@ -0,0 +1,173 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Training CLI for Dynamic Transformers.
+
+This module provides a command-line interface for training models.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="Train a Dynamic Transformer model")
+
+    parser.add_argument(
+        "--config", type=str, required=True, help="Path to configuration file (JSON)"
+    )
+
+    parser.add_argument("--train_data", type=str, required=True, help="Path to training data file")
+
+    parser.add_argument("--val_data", type=str, help="Path to validation data file (optional)")
+
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./experiments",
+        help="Directory to save model checkpoints",
+    )
+
+    parser.add_argument(
+        "--device",
+        type=str,
+        choices=["cuda", "cpu"],
+        default=None,
+        help="Device to use for training",
+    )
+
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
+
+    return parser.parse_args()
+
+
+def main():
+    """Main training function."""
+    args = parse_args()
+
+    print("=" * 60)
+    print("Dynamic Transformers - Training CLI")
+    print("=" * 60)
+
+    # Load configuration
+    print(f"\nLoading configuration from: {args.config}")
+    with open(args.config, "r") as f:
+        config_dict = json.load(f)
+
+    # Load data
+    print(f"\nLoading training data from: {args.train_data}")
+    import pandas as pd
+
+    train_df = (
+        pd.read_csv(args.train_data)
+        if args.train_data.endswith(".csv")
+        else pd.read_json(args.train_data)
+    )
+
+    val_df = None
+    if args.val_data:
+        print(f"Loading validation data from: {args.val_data}")
+        val_df = (
+            pd.read_csv(args.val_data)
+            if args.val_data.endswith(".csv")
+            else pd.read_json(args.val_data)
+        )
+
+    # Import library components
+    from dytr import (
+        DynamicTransformer,
+        ModelConfig,
+        SingleDatasetProcessing,
+        TaskConfig,
+        Trainer,
+        TrainingStrategy,
+        set_seed,
+    )
+
+    # Set seed
+    set_seed(args.seed)
+
+    # Create configuration
+    config = ModelConfig(**config_dict.get("model", {}))
+    if args.device:
+        config.device = args.device
+
+    # Create task configuration
+    task_config = TaskConfig(
+        task_name=config_dict.get("task_name", "task"),
+        training_strategy=TrainingStrategy(config_dict.get("strategy", "sentence_classification")),
+        datasets=[{"train": train_df}],
+        num_labels=config_dict.get("num_labels"),
+        text_column=config_dict.get("text_column", "text"),
+        label_column=config_dict.get("label_column", "label"),
+        max_length=config_dict.get("max_length", 256),
+    )
+
+    # Initialize model
+    print("\nInitializing model...")
+    model = DynamicTransformer(config)
+
+    # Prepare datasets
+    print("Preparing datasets...")
+    train_dataset = SingleDatasetProcessing(
+        df=train_df,
+        tokenizer=model.tokenizer,
+        max_len=task_config.max_length,
+        task_name=task_config.task_name,
+        strategy=task_config.training_strategy,
+        num_labels=task_config.num_labels,
+        text_column=task_config.text_column,
+        label_column=task_config.label_column,
+    )
+
+    train_datasets = {task_config.task_name: (train_dataset, task_config.training_strategy)}
+
+    val_datasets = {}
+    if val_df is not None:
+        val_dataset = SingleDatasetProcessing(
+            df=val_df,
+            tokenizer=model.tokenizer,
+            max_len=task_config.max_length,
+            task_name=task_config.task_name,
+            strategy=task_config.training_strategy,
+            num_labels=task_config.num_labels,
+            text_column=task_config.text_column,
+            label_column=task_config.label_column,
+        )
+        val_datasets = {task_config.task_name: (val_dataset, task_config.training_strategy)}
+
+    # Train
+    print("Starting training...")
+    trainer = Trainer(model, config, exp_dir=args.output_dir)
+    model = trainer.train([task_config], train_datasets, val_datasets)
+
+    # Save final model
+    output_path = Path(args.output_dir) / "final_model.pt"
+    model.save_model(str(output_path))
+    print(f"\nModel saved to: {output_path}")
+
+    print("\n" + "=" * 60)
+    print("Training completed successfully!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

src/dytr/core/__init__.py

@@ -0,0 +1,52 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Core transformer components for Dynamic Transformers.
+
+This module contains the main model classes, configuration, and core building blocks
+for the transformer architecture including attention mechanisms, feedforward networks,
+and the main DynamicTransformer class.
+"""
+
+from dytr.core.config import ModelConfig, TaskConfig, TrainingStrategy
+from dytr.core.decoder import DecoderLayer, TransformerDecoder
+from dytr.core.encoder import EncoderLayer, TransformerEncoder
+from dytr.core.exporter import ModelExporter
+from dytr.core.layers import (
+    FeedForward,
+    MultiHeadAttention,
+    RotaryEmbedding,
+    apply_rotary_pos_emb,
+    rotate_half,
+)
+from dytr.core.model import DynamicTransformer
+
+__all__ = [
+    "DynamicTransformer",
+    "ModelConfig",
+    "TaskConfig",
+    "TrainingStrategy",
+    "RotaryEmbedding",
+    "MultiHeadAttention",
+    "FeedForward",
+    "rotate_half",
+    "apply_rotary_pos_emb",
+    "TransformerEncoder",
+    "EncoderLayer",
+    "TransformerDecoder",
+    "DecoderLayer",
+    "ModelExporter",
+]
+

src/dytr/core/config.py

@@ -0,0 +1,196 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Configuration classes for Dynamic Transformers.
+
+This module defines the configuration dataclasses for model architecture,
+training parameters, and task-specific settings.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+
+class TrainingStrategy(Enum):
+    """Enumeration of supported training strategies."""
+
+    CAUSAL_LM = "causal_lm"
+    SEQ2SEQ = "seq2seq"
+    SENTENCE_CLASSIFICATION = "sentence_classification"
+    TOKEN_CLASSIFICATION = "token_classification"
+
+
+@dataclass
+class ModelConfig:
+    """
+    Configuration class for the Dynamic Transformer model.
+
+    Contains all hyperparameters and settings for model architecture,
+    training, and continual learning features.
+    """
+
+    # Architecture parameters
+    embed_dim: int = 256
+    num_layers: int = 6
+    num_heads: int = 8
+    head_dim: int = 256//8 # embed_dim//num_heads
+    ff_mult: int = 4
+    dropout: float = 0.1
+    max_seq_len: int = 256
+
+    # Training parameters
+    learning_rate: float = 3e-4
+    batch_size: int = 16
+    weight_decay: float = 0.01
+    gradient_clip: float = 1.0
+    warmup_steps: int = 1000
+    max_learning_rate: float = 5e-4
+    min_learning_rate: float = 1e-6
+    adam_epsilon: float = 1e-8
+    label_smoothing: float = 0.1
+
+    # Advanced training features
+    fp16: bool = False
+    gradient_accumulation_steps: int = 1
+    max_grad_norm: float = 1.0
+
+    # Evaluation and checkpointing
+    patience: int = 3
+    evaluation_strategy: str = "steps"
+    logging_steps: int = 50
+    validation_check_interval: int = 500
+    load_best_model_at_end: bool = True
+    metric_for_best_model: str = "loss"
+    early_stopping_patience: int = 10
+
+    # Training duration
+    max_train_steps: int = 100000
+    num_train_epochs: int = 3
+    lr_scheduler_type: str = "cosine"
+
+    # Data loading
+    per_device_train_batch_size: int = 8
+    per_device_eval_batch_size: int = 8
+    dataloader_num_workers: int = 2
+    dataloader_pin_memory: bool = True
+
+    # Randomness
+    seed: int = 42
+
+    # Model enhancements
+    task_specific_lr: Dict[str, float] = field(default_factory=dict)
+    task_weights: Dict[str, float] = field(default_factory=dict)
+    use_rotary_embedding: bool = True
+    use_flash_attention: bool = False
+    gradient_checkpointing: bool = False
+    training_from_scratch: bool=False
+
+    # Special tokens
+    special_tokens: Dict[str, str] = field(
+        default_factory=lambda: {
+            #"task_sep": "<|tasksep|>",
+            #"doc_sep": "<|docsep|>",
+            #"answer_start": "<|answer|>",
+            #"bos": "<s>",
+            #"eos": "</s>",
+        }
+    )
+
+    # Task configuration
+    window_size: int = 256
+    stride: int = 64
+    tasks: Dict[str, Dict[str, Any]] = field(default_factory=dict)
+
+    # Tokenizer IDs
+    vocab_size: Optional[int] = None
+    tokenizer_name: str = "bert-base-multilingual-cased"
+    add_tab_newline_vocab: bool= False
+    use_simple_tokenizer: bool = True
+    tokenizer_type: str = 'wordpiece'
+    bos_token_id: Optional[int] = None
+    eos_token_id: Optional[int] = None
+
+    # Continual learning
+    adapter_bottleneck: int = 64
+    use_task_adapters: bool = True
+    ewc_lambda: float = 1000.0
+    replay_buffer_size: int = 1000
+    use_ewc: bool = False
+    use_replay: bool = False
+
+    # Task-specific learning rate multipliers
+    causal_lm_window_size: int = 256
+    causal_lm_stride: int = 128
+    head_lr_mult: float = 1.0 # Learning rate multiplier for task heads
+    decoder_lr_mult: float = 1.0  # Learning rate multiplier for task decoders
+    shared_lr_mult: float = 0.5 # Learning rate multiplier for shared components
+    
+
+    # Device
+    device: str = "cuda" if __import__("torch").cuda.is_available() else "cpu"
+
+
+@dataclass
+class TaskConfig:
+    """
+    Configuration class for individual tasks.
+
+    Defines the dataset, training strategy, and task-specific parameters
+    for a single task in multi-task learning.
+    """
+
+    # Basic task info
+    task_name: str
+    training_strategy: TrainingStrategy
+
+    # Dataset configuration
+    datasets: List[Dict[str, Any]] = None
+    num_labels: Optional[int] = None
+
+    # Data processing
+    max_length: int = 256
+    train_split: str = "train"
+    validation_split: str = "validation"
+    test_split: str = "test"
+    label_to_ids: Optional[Dict[str,int]]= None
+
+    # Column names
+    text_column: str = "text"
+    label_column: str = "label"
+    source_column: str = "source"
+    target_column: str = "target"
+
+    # Sampling
+    sample_size: Optional[int] = None
+    validation_sample_size: int = 100
+    test_sample_size: int = 100
+
+    # Task-specific options
+    prompt_template: Optional[str] = None
+    metrics: List[str] = field(default_factory=lambda: ["loss", "accuracy"])
+    shuffle_data: bool = True
+    min_text_length: int = 10
+    balance_classes: bool = False
+    dataset_config: Optional[Dict[str, Any]] = None
+    is_generation: bool = False
+
+    # QA specific
+    question_column: Optional[str] = None
+    answer_column: Optional[str] = None
+    conversations_column: Optional[str] = None

src/dytr/core/decoder.py

@@ -0,0 +1,300 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Transformer decoder implementation for Dynamic Transformers.
+
+This module contains the decoder layers and the complete transformer decoder
+with support for cross-attention, causal masking, and text generation.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dytr.core.layers import FeedForward, MultiHeadAttention
+
+
+class DecoderLayer(nn.Module):
+    """
+    Single transformer decoder layer with self-attention, cross-attention, and feedforward network.
+
+    Args:
+        layer_id: Layer index
+        embed_dim: Embedding dimension
+        config: Model configuration
+        has_cross_attention: Whether to include cross-attention (for seq2seq)
+    """
+
+    def __init__(self, layer_id: int, embed_dim: int, config, has_cross_attention: bool = True):
+        super().__init__()
+        self.layer_id = layer_id
+        self.embed_dim = embed_dim
+        self.config = config
+        self.has_cross_attention = has_cross_attention
+
+        # Self-attention (causal)
+        self.self_attention = MultiHeadAttention(
+            embed_dim, config.num_heads, config.head_dim, config, layer_id, is_cross_attention=False
+        )
+
+        # Cross-attention (if enabled)
+        if has_cross_attention:
+            self.cross_attention = MultiHeadAttention(
+                embed_dim,
+                config.num_heads,
+                config.head_dim,
+                config,
+                layer_id,
+                is_cross_attention=True,
+            )
+
+        # Feedforward network
+        self.ffn = FeedForward(embed_dim, config.ff_mult, config, layer_id)
+
+        # Layer normalizations
+        self.norm1 = nn.LayerNorm(embed_dim)
+        if has_cross_attention:
+            self.norm2 = nn.LayerNorm(embed_dim)
+        self.norm3 = nn.LayerNorm(embed_dim)
+
+        # Task adapters
+        if config.use_task_adapters:
+            self.self_attn_adapter = nn.Sequential(
+                nn.Linear(embed_dim, config.adapter_bottleneck),
+                nn.GELU(),
+                nn.Linear(config.adapter_bottleneck, embed_dim),
+            )
+            if has_cross_attention:
+                self.cross_attn_adapter = nn.Sequential(
+                    nn.Linear(embed_dim, config.adapter_bottleneck),
+                    nn.GELU(),
+                    nn.Linear(config.adapter_bottleneck, embed_dim),
+                )
+            self.ffn_adapter = nn.Sequential(
+                nn.Linear(embed_dim, config.adapter_bottleneck),
+                nn.GELU(),
+                nn.Linear(config.adapter_bottleneck, embed_dim),
+            )
+        else:
+            self.self_attn_adapter = None
+            self.cross_attn_adapter = None
+            self.ffn_adapter = None
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        encoder_output: Optional[torch.Tensor] = None,
+        src_mask: Optional[torch.Tensor] = None,
+        tgt_mask: Optional[torch.Tensor] = None,
+        task_name: Optional[str] = None,
+    ):
+        """
+        Forward pass through decoder layer.
+
+        Args:
+            x: Input tensor [batch_size, seq_len, embed_dim]
+            encoder_output: Encoder output for cross-attention
+            src_mask: Source attention mask
+            tgt_mask: Target attention mask (causal)
+            task_name: Current task name for adapter selection
+        """
+        # Self-attention with residual connection
+        residual = x
+        x = self.norm1(x)
+        x = self.self_attention(x, encoder_output=None, mask=tgt_mask, is_causal=True)
+        if self.self_attn_adapter is not None and task_name is not None:
+            x = x + self.self_attn_adapter(x)
+        x = residual + x
+
+        # Cross-attention with residual connection (if enabled)
+        if self.has_cross_attention and encoder_output is not None:
+            residual = x
+            x = self.norm2(x)
+            x = self.cross_attention(
+                x, encoder_output=encoder_output, mask=src_mask, is_causal=False
+            )
+            if self.cross_attn_adapter is not None and task_name is not None:
+                x = x + self.cross_attn_adapter(x)
+            x = residual + x
+
+        # Feedforward with residual connection
+        residual = x
+        x = self.norm3(x)
+        x = self.ffn(x)
+        if self.ffn_adapter is not None and task_name is not None:
+            x = x + self.ffn_adapter(x)
+        x = residual + x
+
+        return x
+
+
+class TransformerDecoder(nn.Module):
+    """
+    Complete transformer decoder consisting of multiple decoder layers.
+
+    Supports both causal LM (no cross-attention) and seq2seq (with cross-attention) modes.
+
+    Args:
+        config: Model configuration
+        task_name: Name of the task this decoder belongs to
+        has_cross_attention: Whether to include cross-attention layers
+    """
+
+    def __init__(self, config, task_name: str, has_cross_attention: bool = True):
+        super().__init__()
+        self.config = config
+        self.task_name = task_name
+        self.has_cross_attention = has_cross_attention
+
+        # Embedding will be set later to share with encoder
+        self.embedding = None
+
+        # Decoder layers
+        self.layers = nn.ModuleList(
+            [
+                DecoderLayer(i, config.embed_dim, config, has_cross_attention)
+                for i in range(config.num_layers)
+            ]
+        )
+
+        # Final layer norm and output projection
+        self.norm = nn.LayerNorm(config.embed_dim)
+        self.output_proj = nn.Linear(config.embed_dim, config.vocab_size)
+
+    def set_shared_embedding(self, embedding: nn.Embedding):
+        """Share embedding weights with encoder."""
+        self.embedding = embedding
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        encoder_output: Optional[torch.Tensor] = None,
+        src_mask: Optional[torch.Tensor] = None,
+        tgt_mask: Optional[torch.Tensor] = None,
+        task_name: Optional[str] = None,
+    ):
+        """
+        Forward pass through transformer decoder.
+
+        Args:
+            input_ids: Input token IDs [batch_size, seq_len]
+            encoder_output: Encoder output for cross-attention
+            src_mask: Source attention mask
+            tgt_mask: Target attention mask
+            task_name: Current task name for adapter selection
+        """
+        x = self.embedding(input_ids)
+
+        for layer in self.layers:
+            if self.has_cross_attention:
+                x = layer(x, encoder_output, src_mask, tgt_mask, task_name or self.task_name)
+            else:
+                x = layer(x, None, None, tgt_mask, task_name or self.task_name)
+
+        x = self.norm(x)
+        logits = self.output_proj(x)
+        return logits
+
+    def generate(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        encoder_output: Optional[torch.Tensor] = None,
+        src_mask: Optional[torch.Tensor] = None,
+        max_len: int = 50,
+        temperature: float = 1.0,
+        top_k: int = 0,
+        top_p: float = 0.0,
+        repetition_penalty: float = 1.0,
+    ):
+        """
+        Generate text using the decoder.
+
+        Args:
+            input_ids: Optional input token IDs for causal LM
+            encoder_output: Encoder output for seq2seq
+            src_mask: Source attention mask
+            max_len: Maximum generation length
+            temperature: Sampling temperature (higher = more random)
+            top_k: Top-k sampling (0 to disable)
+            top_p: Top-p (nucleus) sampling (0 to disable)
+            repetition_penalty: Penalty for repeating tokens
+        """
+        batch_size = encoder_output.size(0) if encoder_output is not None else input_ids.size(0)
+        device = encoder_output.device if encoder_output is not None else input_ids.device
+
+        bos_id = self.config.bos_token_id or 101  # Default BOS token
+        eos_id = self.config.eos_token_id or 102  # Default EOS token
+
+        # Initialize generated sequence
+        if input_ids is not None:
+            generated = input_ids.clone()
+        else:
+            generated = torch.ones(batch_size, 1, dtype=torch.long, device=device) * bos_id
+
+        # Generate tokens autoregressively
+        for _ in range(max_len):
+            # Create causal mask for current sequence
+            tgt_mask = torch.ones(batch_size, generated.size(1), device=device)
+
+            # Forward pass
+            if self.has_cross_attention:
+                logits = self.forward(generated, encoder_output, src_mask, tgt_mask, self.task_name)
+            else:
+                logits = self.forward(generated, tgt_mask=tgt_mask, task_name=self.task_name)
+
+            # Get next token logits and apply temperature
+            next_token_logits = logits[:, -1, :] / temperature
+
+            # Apply repetition penalty
+            for token_id in set(generated[0].tolist()):
+                next_token_logits[:, token_id] /= repetition_penalty
+
+            # Top-k filtering
+            if top_k > 0:
+                indices_to_remove = (
+                    next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                )
+                next_token_logits[indices_to_remove] = float("-inf")
+
+            # Top-p (nucleus) filtering
+            if top_p > 0.0 and top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+
+                for batch_idx in range(next_token_logits.size(0)):
+                    indices_to_remove = sorted_indices[batch_idx][
+                        sorted_indices_to_remove[batch_idx]
+                    ]
+                    next_token_logits[batch_idx, indices_to_remove] = float("-inf")
+
+            # Sample next token
+            probs = F.softmax(next_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            generated = torch.cat([generated, next_token], dim=-1)
+
+            # Stop if all sequences have generated EOS
+            if (next_token == eos_id).all():
+                break
+
+        return generated

src/dytr/core/encoder.py

@@ -0,0 +1,143 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Transformer encoder implementation for Dynamic Transformers.
+
+This module contains the encoder layers and the complete transformer encoder
+with support for task adapters.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from dytr.core.layers import FeedForward, MultiHeadAttention
+
+
+class EncoderLayer(nn.Module):
+    """
+    Single transformer encoder layer with self-attention and feedforward network.
+
+    Args:
+        layer_id: Layer index
+        embed_dim: Embedding dimension
+        config: Model configuration
+    """
+
+    def __init__(self, layer_id: int, embed_dim: int, config):
+        super().__init__()
+        self.layer_id = layer_id
+        self.embed_dim = embed_dim
+        self.config = config
+
+        # Self-attention components
+        self.attention = MultiHeadAttention(
+            embed_dim, config.num_heads, config.head_dim, config, layer_id, is_cross_attention=False
+        )
+        self.attention_norm = nn.LayerNorm(embed_dim)
+
+        # Feedforward components
+        self.ffn_norm = nn.LayerNorm(embed_dim)
+        self.ffn = FeedForward(embed_dim, config.ff_mult, config, layer_id)
+
+        # Task adapter
+        if config.use_task_adapters:
+            self.adapter = nn.Sequential(
+                nn.Linear(embed_dim, config.adapter_bottleneck),
+                nn.GELU(),
+                nn.Linear(config.adapter_bottleneck, embed_dim),
+            )
+        else:
+            self.adapter = None
+
+    def forward(
+        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, task_name: Optional[str] = None
+    ):
+        """
+        Forward pass through encoder layer.
+
+        Args:
+            x: Input tensor [batch_size, seq_len, embed_dim]
+            mask: Attention mask
+            task_name: Current task name for adapter selection
+        """
+        # Self-attention with residual connection
+        residual = x
+        x = self.attention_norm(x)
+        x = self.attention(x, encoder_output=None, mask=mask, is_causal=False)
+        x = residual + x
+
+        # Apply task adapter if available
+        if self.adapter is not None and task_name is not None:
+            x = x + self.adapter(x)
+
+        # Feedforward with residual connection
+        residual = x
+        x = self.ffn_norm(x)
+        x = self.ffn(x)
+        x = residual + x
+
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    """
+    Complete transformer encoder consisting of multiple encoder layers.
+
+    Args:
+        config: Model configuration
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        # Token embedding
+        self.embedding = nn.Embedding(config.vocab_size, config.embed_dim, padding_idx=0)
+        nn.init.normal_(self.embedding.weight, mean=0.0, std=0.02)
+
+        # Encoder layers
+        self.layers = nn.ModuleList(
+            [EncoderLayer(i, config.embed_dim, config) for i in range(config.num_layers)]
+        )
+
+        # Final layer norm
+        self.norm = nn.LayerNorm(config.embed_dim)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        task_name: Optional[str] = None,
+    ):
+        """
+        Forward pass through transformer encoder.
+
+        Args:
+            input_ids: Input token IDs [batch_size, seq_len]
+            mask: Attention mask
+            task_name: Current task name for adapter selection
+        """
+        x = self.embedding(input_ids)
+
+        for layer in self.layers:
+            x = layer(x, mask, task_name)
+
+        x = self.norm(x)
+        return x

src/dytr/core/exporter.py

@@ -0,0 +1,356 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Model exporter for extracting single-task models from multi-task DynamicTransformer.
+
+This module provides functionality to export individual task models for inference,
+supporting all training strategies.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from pathlib import Path
+from typing import Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dytr.core.config import TrainingStrategy
+
+
+class ModelExporter:
+    """
+    Exports single-task models from a multi-task DynamicTransformer.
+
+    This class allows extracting individual task models for deployment,
+    reducing memory footprint and simplifying inference.
+
+    Args:
+        model: The DynamicTransformer instance to export from
+    """
+
+    def __init__(self, model):
+        self.model = model
+        self.tokenizer = model.tokenizer
+        self.config = model.config
+
+    def export_single_task(self, task_name: str, save_path: str):
+        """
+        Export a single task model to a file.
+
+        Args:
+            task_name: Name of the task to export
+            save_path: Path where to save the exported model
+        """
+        if task_name not in self.model.current_tasks:
+            raise ValueError(f"Task {task_name} not found in model")
+
+        strategy = TrainingStrategy(self.model.current_tasks[task_name]["strategy"])
+        num_labels = self.model.current_tasks[task_name].get("num_labels")
+
+        # Sentence Classification Task
+        if strategy == TrainingStrategy.SENTENCE_CLASSIFICATION:
+            if num_labels is None:
+                task_head = self.model.task_heads[task_name]
+                num_labels = (
+                    task_head[-1].out_features
+                    if isinstance(task_head, nn.Sequential)
+                    else task_head.out_features
+                )
+
+            class SingleTaskModel(nn.Module):
+                def __init__(self, encoder, task_head, config, tokenizer, num_labels, task_name):
+                    super().__init__()
+                    self.encoder = encoder
+                    self.task_head = task_head
+                    self.config = config
+                    self.tokenizer = tokenizer
+                    self.num_labels = num_labels
+                    self.task_name = task_name
+                    self.strategy = "sentence_classification"
+
+                def forward(self, input_ids, attention_mask=None):
+                    x = self.encoder(input_ids, attention_mask, self.task_name)
+                    if attention_mask is not None:
+                        mask_sum = attention_mask.sum(dim=1, keepdim=True).clamp(min=1)
+                        pooled = (x * attention_mask.unsqueeze(-1)).sum(dim=1) / mask_sum
+                    else:
+                        pooled = x.mean(dim=1)
+                    logits = self.task_head(pooled)
+                    return logits
+
+                def generate(self, text, **kwargs):
+                    self.eval()
+                    inputs = self.tokenizer(text, return_tensors="pt", add_special_tokens=False)
+                    inputs = {k: v.to(self.config.device) for k, v in inputs.items()}
+                    with torch.no_grad():
+                        logits = self.forward(inputs["input_ids"], inputs.get("attention_mask"))
+                        probs = F.softmax(logits, dim=-1)
+                        pred_id = probs.argmax().item()
+                        return {
+                            "text": text,
+                            "prediction": pred_id,
+                            "probabilities": probs[0].cpu().tolist(),
+                            "logits": logits[0].cpu().tolist(),
+                            "strategy": self.strategy,
+                        }
+
+            model = SingleTaskModel(
+                self.model.encoder,
+                self.model.task_heads[task_name],
+                self.config,
+                self.tokenizer,
+                num_labels,
+                task_name,
+            )
+
+            torch.save(
+                {
+                    "model_state": model.state_dict(),
+                    "config": self.config,
+                    "task_name": task_name,
+                    "strategy": strategy.value,
+                    "tokenizer_name": self.config.tokenizer_name,
+                    "num_labels": num_labels,
+                },
+                save_path,
+            )
+            return model
+
+        # Token Classification Task
+        elif strategy == TrainingStrategy.TOKEN_CLASSIFICATION:
+            if num_labels is None:
+                task_head = self.model.task_heads[task_name]
+                num_labels = (
+                    task_head[-1].out_features
+                    if isinstance(task_head, nn.Sequential)
+                    else task_head.out_features
+                )
+
+            class SingleTaskModel(nn.Module):
+                def __init__(self, encoder, task_head, config, tokenizer, num_labels, task_name):
+                    super().__init__()
+                    self.encoder = encoder
+                    self.task_head = task_head
+                    self.config = config
+                    self.tokenizer = tokenizer
+                    self.num_labels = num_labels
+                    self.task_name = task_name
+                    self.strategy = "token_classification"
+
+                def forward(self, input_ids, attention_mask=None):
+                    x = self.encoder(input_ids, attention_mask, self.task_name)
+                    logits = self.task_head(x)
+                    return logits
+
+                def generate(self, text, **kwargs):
+                    self.eval()
+                    inputs = self.tokenizer(text, return_tensors="pt", add_special_tokens=False)
+                    inputs = {k: v.to(self.config.device) for k, v in inputs.items()}
+                    with torch.no_grad():
+                        logits = self.forward(inputs["input_ids"], inputs.get("attention_mask"))
+                        predictions = logits.argmax(dim=-1)[0].cpu().tolist()
+                        tokens = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+                        return {
+                            "text": text,
+                            "tokens": tokens,
+                            "predictions": predictions,
+                            "pairs": list(zip(tokens, predictions)),
+                            "strategy": self.strategy,
+                        }
+
+            model = SingleTaskModel(
+                self.model.encoder,
+                self.model.task_heads[task_name],
+                self.config,
+                self.tokenizer,
+                num_labels,
+                task_name,
+            )
+
+            torch.save(
+                {
+                    "model_state": model.state_dict(),
+                    "config": self.config,
+                    "task_name": task_name,
+                    "strategy": strategy.value,
+                    "tokenizer_name": self.config.tokenizer_name,
+                    "num_labels": num_labels,
+                },
+                save_path,
+            )
+            return model
+
+        # Seq2Seq Task
+        elif strategy == TrainingStrategy.SEQ2SEQ:
+
+            class Seq2SeqModel(nn.Module):
+                def __init__(self, encoder, decoder, config, tokenizer, task_name):
+                    super().__init__()
+                    self.encoder = encoder
+                    self.decoder = decoder
+                    self.config = config
+                    self.tokenizer = tokenizer
+                    self.task_name = task_name
+                    self.strategy = "seq2seq"
+
+                def forward(self, input_ids, attention_mask=None, labels=None):
+                    encoder_output = self.encoder(input_ids, attention_mask, self.task_name)
+                    if labels is not None:
+                        decoder_input = labels[:, :-1]
+                        logits = self.decoder(decoder_input, encoder_output, attention_mask)
+                        return logits
+                    return encoder_output
+
+                def generate(
+                    self,
+                    text,
+                    max_length=50,
+                    temperature=0.8,
+                    top_p=0.9,
+                    top_k=40,
+                    repetition_penalty=1.1,
+                    **kwargs,
+                ):
+                    self.eval()
+                    inputs = self.tokenizer(text, return_tensors="pt", add_special_tokens=False)
+                    inputs = {k: v.to(self.config.device) for k, v in inputs.items()}
+                    with torch.no_grad():
+                        encoder_output = self.encoder(
+                            inputs["input_ids"], inputs["attention_mask"], self.task_name
+                        )
+                        generated_ids = self.decoder.generate(
+                            encoder_output=encoder_output,
+                            src_mask=inputs["attention_mask"],
+                            max_len=max_length,
+                            temperature=temperature,
+                            top_k=top_k,
+                            top_p=top_p,
+                            repetition_penalty=repetition_penalty,
+                        )
+                        generated_text = self.tokenizer.decode(
+                            generated_ids[0], skip_special_tokens=True
+                        )
+                        return {
+                            "source": text,
+                            "generated": generated_text,
+                            "strategy": self.strategy,
+                        }
+
+            model = Seq2SeqModel(
+                self.model.encoder,
+                self.model.decoders[task_name],
+                self.config,
+                self.tokenizer,
+                task_name,
+            )
+
+            torch.save(
+                {
+                    "model_state": model.state_dict(),
+                    "config": self.config,
+                    "task_name": task_name,
+                    "strategy": strategy.value,
+                    "tokenizer_name": self.config.tokenizer_name,
+                },
+                save_path,
+            )
+            return model
+
+        # Causal LM Task
+        elif strategy == TrainingStrategy.CAUSAL_LM:
+
+            class DecoderOnlyModel(nn.Module):
+                def __init__(self, decoder, config, tokenizer, task_name):
+                    super().__init__()
+                    self.decoder = decoder
+                    self.config = config
+                    self.tokenizer = tokenizer
+                    self.task_name = task_name
+                    self.strategy = "causal_lm"
+
+                def forward(self, input_ids, attention_mask=None):
+                    return self.decoder(input_ids)
+
+                def generate(
+                    self,
+                    prompt,
+                    max_length=50,
+                    temperature=0.8,
+                    top_p=0.9,
+                    top_k=40,
+                    repetition_penalty=1.1,
+                    **kwargs,
+                ):
+                    self.eval()
+                    inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+                    inputs = {k: v.to(self.config.device) for k, v in inputs.items()}
+                    with torch.no_grad():
+                        generated_ids = self.decoder.generate(
+                            input_ids=inputs["input_ids"],
+                            max_len=max_length,
+                            temperature=temperature,
+                            top_k=top_k,
+                            top_p=top_p,
+                            repetition_penalty=repetition_penalty,
+                        )
+                        generated_text = self.tokenizer.decode(
+                            generated_ids[0], skip_special_tokens=True
+                        )
+                        return {
+                            "prompt": prompt,
+                            "generated": generated_text,
+                            "strategy": self.strategy,
+                        }
+
+            model = DecoderOnlyModel(
+                self.model.decoders[task_name], self.config, self.tokenizer, task_name
+            )
+
+            torch.save(
+                {
+                    "model_state": model.state_dict(),
+                    "config": self.config,
+                    "task_name": task_name,
+                    "strategy": strategy.value,
+                    "tokenizer_name": self.config.tokenizer_name,
+                },
+                save_path,
+            )
+            return model
+
+    def export_all_tasks(self, save_dir: str) -> Dict[str, nn.Module]:
+        """
+        Export all tasks from the model.
+
+        Args:
+            save_dir: Directory where to save exported models
+
+        Returns:
+            Dictionary mapping task names to exported models
+        """
+        save_dir = Path(save_dir)
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+        exported_models = {}
+        for task_name in self.model.current_tasks.keys():
+            save_path = save_dir / f"{task_name}.pt"
+            print(f"  Exporting {task_name} -> {save_path}")
+            model = self.export_single_task(task_name, str(save_path))
+            exported_models[task_name] = model
+
+        return exported_models

src/dytr/core/layers.py

@@ -0,0 +1,247 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Core transformer layers including attention mechanisms and feedforward networks.
+
+This module implements the fundamental building blocks of the transformer architecture:
+multi-head attention with rotary positional embeddings, feedforward networks,
+and supporting utility functions.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RotaryEmbedding(nn.Module):
+    """
+    Rotary Position Embedding (RoPE) for transformer models.
+
+    Implements rotary positional embeddings that encode position information
+    by rotating query and key vectors in the complex plane.
+
+    Args:
+        dim: Dimension of the embedding
+        max_seq_len: Maximum sequence length
+        base: Base for the frequency computation
+    """
+
+    def __init__(self, dim: int, max_seq_len: int = 2048, base: int = 10000):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self._build_cache(max_seq_len)
+
+    def _build_cache(self, seq_len: int):
+        """Build and cache cosine and sine values for rotary embeddings."""
+        t = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos())
+        self.register_buffer("sin_cached", emb.sin())
+
+    def forward(self, x: torch.Tensor, seq_len: int):
+        """Get cosine and sine values for rotary embeddings."""
+        if seq_len > self.max_seq_len:
+            self._build_cache(seq_len)
+        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate half of the hidden dimensions."""
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    """Apply rotary position embeddings to query and key tensors."""
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention mechanism with support for cross-attention and rotary embeddings.
+
+    Args:
+        embed_dim: Embedding dimension
+        num_heads: Number of attention heads
+        head_dim: Dimension of each attention head
+        config: Model configuration
+        layer_id: Layer index for debugging
+        is_cross_attention: Whether this is cross-attention layer
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        head_dim: int,
+        config,
+        layer_id: int,
+        is_cross_attention: bool = False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.total_dim = num_heads * head_dim
+        self.scale = head_dim**-0.5
+        self.config = config
+        self.layer_id = layer_id
+        self.is_cross_attention = is_cross_attention
+
+        self.q_proj = nn.Linear(embed_dim, self.total_dim)
+        self.k_proj = nn.Linear(embed_dim, self.total_dim)
+        self.v_proj = nn.Linear(embed_dim, self.total_dim)
+        self.out_proj = nn.Linear(self.total_dim, embed_dim)
+        self.dropout = nn.Dropout(config.dropout)
+
+        if config.use_rotary_embedding and not is_cross_attention:
+            self.rotary = RotaryEmbedding(head_dim, config.max_seq_len)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        encoder_output: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        is_causal: bool = False,
+    ):
+        """
+        Forward pass for multi-head attention.
+
+        Args:
+            x: Input tensor [batch_size, seq_len, embed_dim]
+            encoder_output: Encoder output for cross-attention
+            mask: Attention mask
+            is_causal: Whether to apply causal masking
+        """
+        batch_size, seq_len, _ = x.shape
+
+        # Project and reshape queries
+        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Handle keys and values (self-attention or cross-attention)
+        if self.is_cross_attention and encoder_output is not None:
+            k = (
+                self.k_proj(encoder_output)
+                .view(batch_size, encoder_output.size(1), self.num_heads, self.head_dim)
+                .transpose(1, 2)
+            )
+            v = (
+                self.v_proj(encoder_output)
+                .view(batch_size, encoder_output.size(1), self.num_heads, self.head_dim)
+                .transpose(1, 2)
+            )
+            kv_seq_len = encoder_output.size(1)
+        else:
+            k = (
+                self.k_proj(x)
+                .view(batch_size, seq_len, self.num_heads, self.head_dim)
+                .transpose(1, 2)
+            )
+            v = (
+                self.v_proj(x)
+                .view(batch_size, seq_len, self.num_heads, self.head_dim)
+                .transpose(1, 2)
+            )
+            kv_seq_len = seq_len
+
+        # Apply rotary embeddings
+        if (
+            self.config.use_rotary_embedding
+            and not self.is_cross_attention
+            and encoder_output is None
+        ):
+            cos, sin = self.rotary(x, seq_len)
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        # Compute attention scores
+        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # Apply masks
+        if mask is not None:
+            if mask.dim() == 2:
+                mask = mask[:, None, None, :]
+            if mask.size(-1) != kv_seq_len:
+                if mask.size(-1) > kv_seq_len:
+                    mask = mask[..., :kv_seq_len]
+                else:
+                    pad_len = kv_seq_len - mask.size(-1)
+                    mask = F.pad(mask, (0, pad_len), value=1)
+            if mask.size(1) != self.num_heads:
+                mask = mask.expand(-1, self.num_heads, -1, -1)
+            attn_scores = attn_scores.masked_fill(mask == 0, float("-inf"))
+
+        # Apply causal mask if needed
+        if is_causal and not self.is_cross_attention:
+            causal_mask = torch.triu(
+                torch.ones(seq_len, kv_seq_len, device=x.device) * float("-inf"), diagonal=1
+            )
+            attn_scores = attn_scores + causal_mask[None, None, :, :]
+
+        # Apply softmax and dropout
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        # Apply attention to values
+        context = torch.matmul(attn_weights, v)
+        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.total_dim)
+        output = self.out_proj(context)
+
+        return output
+
+
+class FeedForward(nn.Module):
+    """
+    FeedForward network with GELU activation and SwiGLU-style gating.
+
+    Implements a transformer feedforward layer with gating mechanism
+    similar to SwiGLU used in modern LLMs.
+
+    Args:
+        embed_dim: Embedding dimension
+        ff_mult: Multiplier for feedforward dimension
+        config: Model configuration
+        layer_id: Layer index for debugging
+    """
+
+    def __init__(self, embed_dim: int, ff_mult: int, config, layer_id: int):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.ff_dim = embed_dim * ff_mult
+        self.config = config
+        self.layer_id = layer_id
+
+        self.gate_proj = nn.Linear(embed_dim, self.ff_dim)
+        self.up_proj = nn.Linear(embed_dim, self.ff_dim)
+        self.down_proj = nn.Linear(self.ff_dim, embed_dim)
+        self.activation = nn.GELU()
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x: torch.Tensor):
+        """Forward pass through feedforward network."""
+        return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))

src/dytr/core/model.py

@@ -0,0 +1,619 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Main DynamicTransformer model implementation.
+
+This module contains the core DynamicTransformer class that manages multiple tasks,
+task-specific heads and decoders, and handles different training strategies.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dytr.core.config import ModelConfig, TaskConfig, TrainingStrategy
+from dytr.core.decoder import TransformerDecoder
+from dytr.core.encoder import TransformerEncoder
+from dytr.core.exporter import ModelExporter
+from dytr.memory.ewc import EWC
+from dytr.memory.replay import ReplayBuffer
+
+from collections import defaultdict
+import numpy as np
+
+
+
+
+def calculate_f_scores(preds, labels, average='macro', beta=1.0):
+    """
+    Calculate F-beta score (F1 when beta=1, F0.5 when beta=0.5)
+    
+    Args:
+        preds: Tensor of predictions
+        labels: Tensor of ground truth labels
+        average: 'macro', 'micro', 'weighted', or None
+        beta: Beta parameter for F-beta score
+    Returns:
+        f_score: F-beta score
+        per_class: Dictionary of per-class scores (if average is None)
+    """
+    
+    
+    if labels.numel() == 0:
+        return 0.0, {}
+    
+    # Convert to CPU numpy for calculation
+    preds_np = preds.cpu().numpy()
+    labels_np = labels.cpu().numpy()
+    
+    # Get unique classes
+    classes = np.unique(np.concatenate([labels_np, preds_np]))
+    n_classes = len(classes)
+    
+    if n_classes == 0:
+        return 0.0, {}
+    
+    # Initialize confusion matrix components
+    tp = defaultdict(int)
+    fp = defaultdict(int)
+    fn = defaultdict(int)
+    
+    for pred, label in zip(preds_np, labels_np):
+        if pred == label:
+            tp[label] += 1
+        else:
+            fp[pred] += 1
+            fn[label] += 1
+    
+    # Calculate precision, recall, and f-beta for each class
+    precision = {}
+    recall = {}
+    f_beta = {}
+    
+    for cls in classes:
+        p = tp[cls] / (tp[cls] + fp[cls]) if (tp[cls] + fp[cls]) > 0 else 0.0
+        r = tp[cls] / (tp[cls] + fn[cls]) if (tp[cls] + fn[cls]) > 0 else 0.0
+        
+        precision[cls] = p
+        recall[cls] = r
+        
+        if p + r > 0:
+            beta_sq = beta * beta
+            f_beta[cls] = (1 + beta_sq) * (p * r) / ((beta_sq * p) + r)
+        else:
+            f_beta[cls] = 0.0
+    
+    if average == 'macro':
+        return sum(f_beta.values()) / n_classes, f_beta
+    elif average == 'weighted':
+        # Count samples per class
+        class_counts = defaultdict(int)
+        for label in labels_np:
+            class_counts[label] += 1
+        total = len(labels_np)
+        weighted_score = sum(f_beta[cls] * (class_counts[cls] / total) for cls in classes)
+        return weighted_score, f_beta
+    elif average == 'micro':
+        # Micro averaging: aggregate all classes
+        total_tp = sum(tp.values())
+        total_fp = sum(fp.values())
+        total_fn = sum(fn.values())
+        micro_p = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0.0
+        micro_r = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0
+        beta_sq = beta * beta
+        micro_f = (1 + beta_sq) * (micro_p * micro_r) / ((beta_sq * micro_p) + micro_r) if (micro_p + micro_r) > 0 else 0.0
+        return micro_f, f_beta
+    else:
+        return f_beta, f_beta
+
+
+
+class DynamicTransformer(nn.Module):
+    """
+    Main Dynamic Transformer model supporting multi-task learning and continual learning.
+
+    This model can handle multiple tasks with different training strategies simultaneously,
+    with support for EWC (Elastic Weight Consolidation) and experience replay to prevent
+    catastrophic forgetting.
+
+    Args:
+        config: Model configuration
+    """
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.current_tasks = {}
+        expected_head_dim = config.embed_dim // config.num_heads
+        if config.embed_dim % config.num_heads != 0:
+            raise ValueError(
+                f"embed_dim ({config.embed_dim}) must be divisible by num_heads ({config.num_heads})"
+            )
+        if config.head_dim != expected_head_dim:
+            print(f"\n⚠️ WARNING: head_dim ({config.head_dim}) does not match expected value ({expected_head_dim})")
+            
+            print(f"   Formula: embed_dim ({config.embed_dim}) // num_heads ({config.num_heads}) = {expected_head_dim}")
+            print(f"   Continuing with user-specified value: {config.head_dim}")
+            print("   (This may cause unexpected behavior)")
+        self.config = config
+
+        # Initialize tokenizer (try HuggingFace first, fallback to SimpleTokenizer)
+        try:
+            from transformers import AutoTokenizer
+            transformers_available = True
+        except:
+            transformers_available = False
+            
+        if transformers_available and not config.use_simple_tokenizer:
+            self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name)
+        else:
+            
+            from dytr.tokenization.download_manager import download_tokenizer_files
+            from dytr.tokenization.simple_tokenizer import SimpleTokenizer
+
+            if "vocab.json" in config.tokenizer_name :
+                import json
+
+                with open(config.tokenizer_name, "r", encoding="utf-8") as f:
+                    vocab = json.load(f)
+                self.tokenizer = SimpleTokenizer(vocab,add_tab_newline_vocab = config.add_tab_newline_vocab, tokenizer_type=config.tokenizer_type)
+            else:
+                try:
+                    vocab, special_token_map = download_tokenizer_files(str(config.tokenizer_name))
+                    self.tokenizer = SimpleTokenizer(vocab, special_token_map,add_tab_newline_vocab = config.add_tab_newline_vocab,tokenizer_type=config.tokenizer_type)
+                except:
+                    self.tokenizer = SimpleTokenizer.from_pretrained(config.tokenizer_name,tokenizer_type=config.tokenizer_type,add_tab_newline_vocab = config.add_tab_newline_vocab)
+
+        # Configure tokenizer
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = (
+                self.tokenizer.eos_token if self.tokenizer.eos_token else "[PAD]"
+            )
+
+        # Add special tokens
+        for token in config.special_tokens.values():
+            if token not in self.tokenizer.get_vocab() and config.training_from_scratch:
+                self.tokenizer.add_tokens([token])
+
+        # Update config with tokenizer info
+        config.vocab_size = len(self.tokenizer)
+        config.bos_token_id = self.tokenizer.bos_token_id or self.tokenizer.cls_token_id
+        config.eos_token_id = self.tokenizer.eos_token_id or self.tokenizer.sep_token_id
+
+        # Shared components
+        self.shared_embedding = nn.Embedding(
+            config.vocab_size, config.embed_dim, padding_idx=self.tokenizer.pad_token_id
+        )
+        nn.init.normal_(self.shared_embedding.weight, mean=0.0, std=0.02)
+
+        # Encoder (shared across all tasks)
+        self.encoder = TransformerEncoder(config)
+        self.encoder.embedding = self.shared_embedding
+
+        # Task-specific components
+        self.decoders = nn.ModuleDict()
+        self.task_heads = nn.ModuleDict()
+
+        # Training state
+        self.training_step = 0
+        self.best_val_loss = float("inf")
+        self.steps_without_improvement = 0
+
+        # Continual learning components
+        self.ewc_penalties = {}
+        self.replay_buffer = ReplayBuffer(config.replay_buffer_size)
+        self.error_mask = None
+
+        # Initialize tasks from config
+        if config.tasks:
+            for task_name, task_info in config.tasks.items():
+                strategy = TrainingStrategy(task_info["strategy"])
+                num_labels = task_info.get("num_labels", config.vocab_size)
+                self._add_task_head(task_name, strategy, num_labels)
+                self.current_tasks[task_name] = task_info
+
+        self.to(config.device)
+
+    def _add_task_head(self, task_name: str, strategy: TrainingStrategy, num_labels: int):
+        """Add task-specific head or decoder."""
+        if strategy == TrainingStrategy.SENTENCE_CLASSIFICATION:
+            head = nn.Sequential(
+                nn.Linear(self.config.embed_dim, self.config.embed_dim // 2),
+                nn.LayerNorm(self.config.embed_dim // 2),
+                nn.GELU(),
+                nn.Dropout(self.config.dropout),
+                nn.Linear(self.config.embed_dim // 2, num_labels),
+            )
+            self.task_heads[task_name] = head.to(self.config.device)
+
+        elif strategy == TrainingStrategy.TOKEN_CLASSIFICATION:
+            head = nn.Sequential(
+                nn.Linear(self.config.embed_dim, self.config.embed_dim),
+                nn.LayerNorm(self.config.embed_dim),
+                nn.GELU(),
+                nn.Dropout(self.config.dropout),
+                nn.Linear(self.config.embed_dim, num_labels),
+            )
+            self.task_heads[task_name] = head.to(self.config.device)
+
+        elif strategy == TrainingStrategy.CAUSAL_LM:
+            decoder = TransformerDecoder(self.config, task_name, has_cross_attention=False)
+            decoder.set_shared_embedding(self.shared_embedding)
+            self.decoders[task_name] = decoder.to(self.config.device)
+
+        elif strategy == TrainingStrategy.SEQ2SEQ:
+            decoder = TransformerDecoder(self.config, task_name, has_cross_attention=True)
+            decoder.set_shared_embedding(self.shared_embedding)
+            self.decoders[task_name] = decoder.to(self.config.device)
+
+    def add_task(self, task_config: TaskConfig):
+        """Add a new task to the model."""
+        task_name = task_config.task_name
+
+        if task_name not in self.task_heads and task_name not in self.decoders:
+            num_labels = task_config.num_labels or self.config.vocab_size
+            self._add_task_head(task_name, task_config.training_strategy, num_labels)
+
+            task_info = {
+                "strategy": task_config.training_strategy.value,
+                "num_labels": num_labels,
+                "best_loss": float("inf"),
+                "best_metrics": {},
+                "steps_without_improvement": 0,
+            }
+            self.current_tasks[task_name] = task_info
+            self.config.tasks[task_name] = task_info
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        task_name: Optional[str] = None,
+        labels: Optional[torch.Tensor] = None,
+        encoder_output: Optional[torch.Tensor] = None,
+        error_mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        Forward pass for the model.
+
+        Handles different training strategies:
+        - SEQ2SEQ: Encoder-decoder with cross-attention
+        - CAUSAL_LM: Decoder-only autoregressive
+        - SENTENCE_CLASSIFICATION: Classification with pooling
+        - TOKEN_CLASSIFICATION: Token-level classification
+        """
+        if error_mask is not None:
+            self.error_mask = error_mask
+
+        if task_name and task_name in self.current_tasks:
+            strategy = TrainingStrategy(self.current_tasks[task_name]["strategy"])
+
+            # Sequence-to-Sequence
+            if strategy == TrainingStrategy.SEQ2SEQ:
+                if encoder_output is None:
+                    encoder_output = self.encoder(input_ids, attention_mask, task_name)
+
+                if task_name not in self.decoders:
+                    raise ValueError(f"No decoder found for task {task_name}")
+
+                decoder = self.decoders[task_name]
+
+                if labels is not None:
+                    decoder_input = labels[:, :-1].clone()
+                    decoder_input[decoder_input == -100] = self.tokenizer.pad_token_id or 0
+                    target_labels = labels[:, 1:]
+                    decoder_mask = (decoder_input != self.tokenizer.pad_token_id).float()
+
+                    logits = decoder(
+                        decoder_input, encoder_output, attention_mask, decoder_mask, task_name
+                    )
+
+                    loss_fn = nn.CrossEntropyLoss(
+                        ignore_index=-100, label_smoothing=self.config.label_smoothing
+                    )
+                    loss = loss_fn(
+                        logits.reshape(-1, self.config.vocab_size), target_labels.reshape(-1)
+                    )
+
+                    metrics = {}
+                    if not self.training:
+                        mask = target_labels != -100
+                        if mask.any():
+                            preds = logits.argmax(dim=-1)
+                            correct = (preds[mask] == target_labels[mask]).float().sum().item()
+                            total = mask.sum().item()
+                            metrics["token_accuracy"] = correct / total if total > 0 else 0.0
+
+                    return {
+                        "logits": logits,
+                        "loss": loss,
+                        "encoder_output": encoder_output,
+                        "metrics": metrics,
+                    }
+                else:
+                    return {"encoder_output": encoder_output}
+
+            # Causal Language Model
+            elif strategy == TrainingStrategy.CAUSAL_LM:
+                decoder = self.decoders[task_name]
+
+                if labels is not None:
+                    decoder_input = labels[:, :-1].clone()
+                    decoder_input[decoder_input == -100] = self.tokenizer.pad_token_id or 0
+                    target_labels = labels[:, 1:]
+
+                    logits = decoder(decoder_input, task_name=task_name)
+                    loss_fn = nn.CrossEntropyLoss(
+                        ignore_index=-100, label_smoothing=self.config.label_smoothing
+                    )
+                    loss = loss_fn(
+                        logits.reshape(-1, self.config.vocab_size), target_labels.reshape(-1)
+                    )
+
+                    metrics = {}
+                    if not self.training:
+                        mask = target_labels != -100
+                        if mask.any():
+                            preds = logits.argmax(dim=-1)
+                            metrics["perplexity"] = torch.exp(loss).item()
+                            correct = (preds[mask] == target_labels[mask]).float().sum().item()
+                            total = mask.sum().item()
+                            metrics["token_accuracy"] = correct / total if total > 0 else 0.0
+
+                    return {"logits": logits, "loss": loss, "metrics": metrics}
+                else:
+                    return {"decoder": decoder}
+
+            # Sentence Classification
+            elif strategy == TrainingStrategy.SENTENCE_CLASSIFICATION:
+                x = self.shared_embedding(input_ids)
+                for layer in self.encoder.layers:
+                    x = layer(x, attention_mask, task_name)
+                x = self.encoder.norm(x)
+
+                # Pooling
+                if attention_mask is not None:
+                    mask_sum = attention_mask.sum(dim=1, keepdim=True).clamp(min=1)
+                    pooled = (x * attention_mask.unsqueeze(-1)).sum(dim=1) / mask_sum
+                else:
+                    pooled = x.mean(dim=1)
+
+                logits = self.task_heads[task_name](pooled)
+
+                if labels is not None:
+                    if labels.dim() > 1:
+                        labels = labels[:, 0]
+                    loss_fn = nn.CrossEntropyLoss(label_smoothing=self.config.label_smoothing)
+                    loss = loss_fn(logits, labels)
+
+                    metrics = {}
+                    if not self.training:
+                        preds = logits.argmax(dim=-1)
+                        correct = (preds == labels).float().sum().item()
+                        total = labels.size(0)
+                        metrics["accuracy"] = correct / total if total > 0 else 0.0
+                        if total>0:
+                            f1_macro, _ = calculate_f_scores(preds, labels, average='macro', beta=1.0)
+                            f1_weighted, _ = calculate_f_scores(preds, labels, average='weighted', beta=1.0)
+                            f1_micro, _ = calculate_f_scores(preds, labels, average='micro', beta=1.0)
+                            metrics["f1_macro"] = f1_macro
+                            metrics["f1_weighted"] = f1_weighted
+                            metrics["f1_micro"] = f1_micro
+
+                    return {"logits": logits, "loss": loss, "metrics": metrics}
+                return {"logits": logits}
+
+            # Token Classification
+            elif strategy == TrainingStrategy.TOKEN_CLASSIFICATION:
+                x = self.shared_embedding(input_ids)
+                for layer in self.encoder.layers:
+                    x = layer(x, attention_mask, task_name)
+                x = self.encoder.norm(x)
+
+                logits = self.task_heads[task_name](x)
+
+                if labels is not None:
+                    loss_fn = nn.CrossEntropyLoss(
+                        ignore_index=-100, label_smoothing=self.config.label_smoothing
+                    )
+                    active = labels.view(-1) != -100
+
+                    if active.any():
+                        logits_flat = logits.view(-1, self.task_heads[task_name][-1].out_features)
+                        labels_flat = labels.view(-1)
+                        loss = loss_fn(logits_flat[active], labels_flat[active])
+                    else:
+                        loss = torch.tensor(0.0, device=logits.device, requires_grad=True)
+
+                    metrics = {}
+                    if not self.training and active.any():
+                        preds = logits.argmax(dim=-1).view(-1)[active]
+                        labels_flat_active = labels_flat[active]
+
+                        correct = (preds == labels_flat_active).float().sum().item()
+                        total = active.sum().item()
+                        metrics["token_accuracy"] = correct / total if total > 0 else 0.0
+                        """
+                        is_error = labels_flat_active != 0
+
+                        if is_error.any():
+                            error_correct = (preds[is_error] == labels_flat_active[is_error]).float().sum().item()
+                            error_total = is_error.sum().item()
+                            metrics["token_accuracy(non-zero tags)"] = error_correct / error_total if error_total > 0 else 0.0
+                        """
+                        if total > 0:
+                            # Calculate macro F1 (averaged over all classes)
+                            f1_macro, per_class_f1 = calculate_f_scores(preds, labels_flat_active, average='macro', beta=1.0)
+                            
+                            #f05_macro, per_class_f05 = calculate_f_scores(preds, labels_flat_active, average='macro', beta=0.5)
+                            
+                            # Calculate weighted F1 (accounting for class imbalance)
+                            f1_weighted, _ = calculate_f_scores(preds, labels_flat_active, average='weighted', beta=1.0)
+                            #f05_weighted, _ = calculate_f_scores(preds, labels_flat_active, average='weighted', beta=0.5)
+                            
+                            # Calculate micro F1 (overall performance)
+                            f1_micro, _ = calculate_f_scores(preds, labels_flat_active, average='micro', beta=1.0)
+                            #f05_micro, _ = calculate_f_scores(preds, labels_flat_active, average='micro', beta=0.5)
+                            
+                            metrics["f1_macro"] = f1_macro
+                            #metrics["f0.5_macro"] = f05_macro
+                            metrics["f1_weighted"] = f1_weighted
+                            #metrics["f0.5_weighted"] = f05_weighted
+                            metrics["f1_micro"] = f1_micro
+                            #metrics["f0.5_micro"] = f05_micro
+                            
+                            # Optional: Add per-class scores (useful for debugging)
+                            # metrics["per_class_f1"] = {str(k): v for k, v in per_class_f1.items()}
+                            
+
+                    return {"logits": logits, "loss": loss, "metrics": metrics}
+                return {"logits": logits}
+
+        # Fallback: just encode
+        x = self.shared_embedding(input_ids)
+        for layer in self.encoder.layers:
+            x = layer(x, attention_mask, None)
+        x = self.encoder.norm(x)
+        return {"hidden_states": x}
+
+    def generate(
+        self,
+        prompt: str,
+        task_name: Optional[str] = None,
+        max_new_tokens: int = 50,
+        temperature: float = 0.8,
+        top_p: float = 0.9,
+        top_k: int = 40,
+        repetition_penalty: float = 1.1,
+        decoding_with_special_tokens: str = "no",
+    ):
+        """Generate text or predictions for a given prompt."""
+        self.eval()
+        inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+        input_ids = inputs["input_ids"].to(self.config.device)
+        attention_mask = inputs["attention_mask"].to(self.config.device)
+
+        if task_name and task_name in self.current_tasks:
+            strategy = TrainingStrategy(self.current_tasks[task_name]["strategy"])
+
+            if strategy == TrainingStrategy.CAUSAL_LM:
+                decoder = self.decoders[task_name]
+                with torch.no_grad():
+                    generated_ids = decoder.generate(
+                        input_ids=input_ids,
+                        max_len=max_new_tokens,
+                        temperature=temperature,
+                        top_k=top_k,
+                        top_p=top_p,
+                        repetition_penalty=repetition_penalty,
+                    )
+                if decoding_with_special_tokens == "no":
+                    return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+                elif decoding_with_special_tokens == "yes":
+                    return self.tokenizer.decode(generated_ids[0], skip_special_tokens=False)
+                else:
+                    return generated_ids
+
+            elif strategy == TrainingStrategy.SEQ2SEQ:
+                with torch.no_grad():
+                    encoder_output = self.encoder(input_ids, attention_mask, task_name)
+                    decoder = self.decoders[task_name]
+                    generated_ids = decoder.generate(
+                        encoder_output=encoder_output,
+                        src_mask=attention_mask,
+                        max_len=max_new_tokens,
+                        temperature=temperature,
+                        top_k=top_k,
+                        top_p=top_p,
+                        repetition_penalty=repetition_penalty,
+                    )
+                if decoding_with_special_tokens == "no":
+                    return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+                elif decoding_with_special_tokens == "yes":
+                    return self.tokenizer.decode(generated_ids[0], skip_special_tokens=False)
+                else:
+                    return generated_ids
+
+            elif strategy == TrainingStrategy.SENTENCE_CLASSIFICATION:
+                with torch.no_grad():
+                    outputs = self.forward(input_ids, attention_mask, task_name=task_name)
+                    logits = outputs["logits"]
+                    probs = F.softmax(logits, dim=-1)
+                    pred_id = probs.argmax().item()
+                    return {
+                        "prediction": pred_id,
+                        "probabilities": probs[0].cpu().tolist(),
+                        "logits": logits[0].cpu().tolist(),
+                    }
+
+            elif strategy == TrainingStrategy.TOKEN_CLASSIFICATION:
+                with torch.no_grad():
+                    outputs = self.forward(input_ids, attention_mask, task_name=task_name)
+                    logits = outputs["logits"]
+                    predictions = logits.argmax(dim=-1)[0].cpu().tolist()
+                    tokens = self.tokenizer.convert_ids_to_tokens(input_ids[0])
+                    return {
+                        "tokens": tokens,
+                        "predictions": predictions,
+                        "pairs": list(zip(tokens, predictions)),
+                    }
+
+        return prompt
+
+    def save_model(self, path: str):
+        """Save the entire model to disk."""
+        self.config.tasks = self.current_tasks
+        state = {
+            "model_state": self.state_dict(),
+            "config": self.config,
+            "current_tasks": self.current_tasks,
+            "training_step": self.training_step,
+            "best_val_loss": self.best_val_loss,
+            "steps_without_improvement": self.steps_without_improvement,
+            "tokenizer_name": self.config.tokenizer_name,
+            "special_tokens": self.config.special_tokens,
+        }
+        torch.save(state, path)
+
+    @classmethod
+    def load_model(cls, path: str, device: Optional[str] = None):
+        """Load a saved model from disk."""
+        try:
+            if not device:
+                device = "cuda" if torch.cuda.is_available() else "cpu"
+
+            state = torch.load(path, map_location=device, weights_only=False)
+            config = state["config"]
+            config.device = device
+
+            model = cls(config)
+            model.load_state_dict(state["model_state"], strict=False)
+            model.current_tasks = state.get("current_tasks", {})
+            model.training_step = state.get("training_step", 0)
+            model.best_val_loss = state.get("best_val_loss", float("inf"))
+            model.steps_without_improvement = state.get("steps_without_improvement", 0)
+
+            return model.to(config.device)
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            return None
+
+    def get_exporter(self):
+        """Get a model exporter instance."""
+        return ModelExporter(self)

src/dytr/memory/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Memory and continual learning modules for Dynamic Transformers.
+
+This module provides implementations of continual learning techniques
+including Elastic Weight Consolidation (EWC) and Experience Replay
+to prevent catastrophic forgetting when training on sequential tasks.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from dytr.memory.ewc import EWC
+from dytr.memory.replay import ReplayBuffer
+
+__all__ = ["EWC", "ReplayBuffer"]

src/dytr/memory/ewc.py

@@ -0,0 +1,112 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Elastic Weight Consolidation (EWC) for continual learning.
+
+This module implements EWC to prevent catastrophic forgetting when training
+on sequential tasks.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from typing import Dict
+
+import torch
+import torch.nn as nn
+
+
+class EWC:
+    """
+    Elastic Weight Consolidation for continual learning.
+
+    EWC computes the Fisher information matrix for important parameters
+    and adds a penalty to prevent them from changing too much when learning
+    new tasks.
+
+    Args:
+        model: The model to apply EWC to
+        task_name: Name of the task this EWC instance is for
+        lambda_param: Regularization strength
+    """
+
+    def __init__(self, model: nn.Module, task_name: str, lambda_param: float = 1000.0):
+        self.model = model
+        self.task_name = task_name
+        self.lambda_param = lambda_param
+        self.params = {
+            n: p.clone().detach() for n, p in model.named_parameters() if p.requires_grad
+        }
+        self.fisher: Dict[str, torch.Tensor] = {}
+
+    def compute_fisher(self, dataloader, device):
+        """
+        Compute Fisher information matrix using the given dataloader.
+
+        Args:
+            dataloader: DataLoader containing samples from the task
+            device: Device to run computation on
+        """
+        # Initialize Fisher matrix
+        for n, p in self.model.named_parameters():
+            if p.requires_grad:
+                self.fisher[n] = torch.zeros_like(p)
+
+        self.model.eval()
+
+        for batch in dataloader:
+            # Only process batches for this task
+            if batch["task_name"][0] != self.task_name:
+                continue
+
+            batch = {
+                k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()
+            }
+            self.model.zero_grad()
+
+            outputs = self.model(
+                input_ids=batch["input_ids"],
+                attention_mask=batch.get("attention_mask"),
+                task_name=self.task_name,
+                labels=batch.get("labels"),
+            )
+
+            loss = outputs["loss"]
+            if loss is not None:
+                loss.backward()
+
+                for n, p in self.model.named_parameters():
+                    if p.requires_grad and p.grad is not None:
+                        self.fisher[n] += p.grad.detach() ** 2
+
+        # Normalize Fisher matrix
+        for n in self.fisher:
+            self.fisher[n] /= len(dataloader)
+
+    def penalty(self, model: nn.Module) -> torch.Tensor:
+        """
+        Compute the EWC penalty for the current model.
+
+        Args:
+            model: Current model state
+
+        Returns:
+            EWC penalty loss
+        """
+        loss = 0.0
+        for n, p in model.named_parameters():
+            if n in self.fisher and n in self.params:
+                loss += (self.fisher[n] * (p - self.params[n]) ** 2).sum()
+        return self.lambda_param * loss

src/dytr/memory/replay.py

@@ -0,0 +1,83 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Experience replay buffer for continual learning.
+
+This module implements a replay buffer that stores samples from previous tasks
+to prevent catastrophic forgetting.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+import random
+from collections import defaultdict
+from typing import Any, List, Optional
+
+
+class ReplayBuffer:
+    """
+    Experience replay buffer for storing samples from previous tasks.
+
+    The buffer maintains a fixed capacity and stores samples along with
+    their task labels for later replay during training.
+
+    Args:
+        capacity: Maximum number of samples to store
+    """
+
+    def __init__(self, capacity: int = 1000):
+        self.capacity = capacity
+        self.buffer = []
+        self.task_indices = defaultdict(list)
+
+    def add_samples(self, task_name: str, samples: List[Any]):
+        """
+        Add samples to the replay buffer.
+
+        Args:
+            task_name: Name of the task these samples belong to
+            samples: List of samples to add
+        """
+        for sample in samples:
+            # Remove oldest sample if buffer is full
+            if len(self.buffer) >= self.capacity:
+                oldest_task = self.buffer[0][0]
+                self.task_indices[oldest_task].pop(0)
+                self.buffer.pop(0)
+
+            self.buffer.append((task_name, sample))
+            self.task_indices[task_name].append(len(self.buffer) - 1)
+
+    def sample(self, batch_size: int, task_name: Optional[str] = None) -> List[Any]:
+        """
+        Sample random samples from the buffer.
+
+        Args:
+            batch_size: Number of samples to retrieve
+            task_name: If specified, only sample from this task
+
+        Returns:
+            List of sampled samples
+        """
+        if task_name and task_name in self.task_indices and self.task_indices[task_name]:
+            indices = random.sample(
+                self.task_indices[task_name], min(batch_size, len(self.task_indices[task_name]))
+            )
+            return [self.buffer[i][1] for i in indices]
+        elif self.buffer:
+            indices = random.sample(range(len(self.buffer)), min(batch_size, len(self.buffer)))
+            return [self.buffer[i][1] for i in indices]#, [self.buffer[i][0] for i in indices]
+        return []

src/dytr/pretrained/__init__.py

@@ -0,0 +1,13 @@
+"""
+Pretrained model loading utilities for dytr.
+
+This module provides functionality to load pretrained models from HuggingFace
+without requiring the transformers library.
+
+Author: Dr. Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+from dytr.pretrained.loader import PretrainedModelLoader
+
+__all__ = ["PretrainedModelLoader"]

src/dytr/pretrained/loader.py

@@ -0,0 +1,867 @@
+"""
+Pretrained model loader for Dynamic Transformer.
+Supports loading encoder-only models (BERT, RoBERTa, DistilBERT, ALBERT) as the shared encoder.
+
+Author: Dr. Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+import os
+import json
+import torch
+import torch.nn as nn
+from pathlib import Path
+from typing import Dict, Optional, Union, List, Tuple
+import requests
+from tqdm import tqdm
+
+from dytr.core.model import DynamicTransformer
+from dytr.core.config import ModelConfig
+
+
+class PretrainedModelLoader:
+    """
+    Load pretrained encoder models from HuggingFace without transformers library.
+    Supports BERT, RoBERTa, DistilBERT, ALBERT architectures.
+    
+    These models are loaded as the shared encoder in DynamicTransformer.
+    
+    Example:
+        >>> from dytr import PretrainedModelLoader, ModelConfig
+        >>> loader = PretrainedModelLoader()
+        >>> config = ModelConfig(tokenizer_name='prajjwal1/bert-tiny', use_simple_tokenizer=False)
+        >>> model = loader.load_pretrained('prajjwal1/bert-tiny', config)
+    """
+
+    # Supported encoder-only model types
+    SUPPORTED_MODELS = ['bert', 'roberta', 'distilbert', 'albert']
+    
+    # Models that are NOT supported (decoder-only or encoder-decoder)
+    UNSUPPORTED_MODELS = ['gpt2', 'gpt', 't5', 'bart', 'pegasus', 'bloom', 'llama']
+
+    def __init__(self, cache_dir: str = "./pretrained_cache"):
+        """
+        Initialize the pretrained model loader.
+        
+        Args:
+            cache_dir: Directory to cache downloaded model files
+        """
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def load_pretrained(self, model_name: str, config: Optional[ModelConfig] = None, 
+                        config_override: Optional[Dict] = None) -> DynamicTransformer:
+        """
+        Load a pretrained encoder model from HuggingFace.
+        
+        Args:
+            model_name: HuggingFace model name (e.g., 'prajjwal1/bert-tiny', 'bert-base-uncased')
+            config: ModelConfig instance (will be created if not provided)
+            config_override: Override configuration parameters
+        
+        Returns:
+            DynamicTransformer model with pretrained encoder weights
+        
+        Raises:
+            ValueError: If model type is not supported (decoder-only models)
+        """
+        # Detect and validate model type
+        model_type = self._detect_model_type(model_name)
+        
+        if model_type in self.UNSUPPORTED_MODELS:
+            raise ValueError(
+                f"\n{'='*60}\n"
+                f"ERROR: Model '{model_name}' is a {model_type.upper()} model which is NOT supported.\n"
+                f"{'='*60}\n"
+                f"DynamicTransformer only supports encoder-only models as the shared encoder.\n\n"
+                f"Supported model types: {self.SUPPORTED_MODELS}\n"
+                f"Examples:\n"
+                f"  - BERT: 'prajjwal1/bert-tiny', 'bert-base-uncased'\n"
+                f"  - RoBERTa: 'roberta-base'\n"
+                f"  - DistilBERT: 'distilbert-base-uncased'\n"
+                f"  - ALBERT: 'albert-base-v2'\n\n"
+                f"Unsupported model types: {self.UNSUPPORTED_MODELS}\n"
+                f"{'='*60}"
+            )
+        
+        if model_type not in self.SUPPORTED_MODELS:
+            raise ValueError(
+                f"Unsupported model type: {model_type}. "
+                f"Supported: {self.SUPPORTED_MODELS}"
+            )
+        
+        if model_type == 'bert':
+            return self.load_bert(model_name, config, config_override)
+        elif model_type == 'roberta':
+            return self.load_roberta(model_name, config, config_override)
+        elif model_type == 'distilbert':
+            return self.load_distilbert(model_name, config, config_override)
+        elif model_type == 'albert':
+            return self.load_albert(model_name, config, config_override)
+        else:
+            raise ValueError(f"Unsupported model type: {model_type}")
+
+    def _detect_model_type(self, model_name: str) -> str:
+        """Detect model type from name or config."""
+        model_name_lower = model_name.lower()
+        
+        # Check by name first
+        if 'xlm-roberta' in model_name_lower:
+            return 'xlm'
+        elif 'roberta' in model_name_lower:
+            return 'roberta'
+        elif 'distilbert' in model_name_lower:
+            return 'distilbert'
+        elif 'albert' in model_name_lower:
+            return 'albert'
+        elif 'bert' in model_name_lower:
+            return 'bert'
+        elif 'gpt' in model_name_lower:
+            return 'gpt2'
+        elif 't5' in model_name_lower:
+            return 't5'
+        elif 'bart' in model_name_lower:
+            return 'bart'
+        elif 'bloom' in model_name_lower:
+            return 'bloom'
+        elif 'llama' in model_name_lower:
+            return 'llama'
+        
+        # Try to detect from config if already downloaded
+        try:
+            config_path = self.cache_dir / model_name / 'config.json'
+            if config_path.exists():
+                with open(config_path, 'r') as f:
+                    config_data = json.load(f)
+                if 'model_type' in config_data:
+                    model_type = config_data['model_type']
+                    if model_type in self.SUPPORTED_MODELS:
+                        return model_type
+                    else:
+                        return model_type
+        except:
+            pass
+        
+        raise ValueError(f"Could not detect model type for {model_name}")
+
+    def download_model_files(self, model_name: str, model_type: str) -> Dict[str, Path]:
+        """
+        Download model files from HuggingFace.
+        
+        Args:
+            model_name: HuggingFace model name
+            model_type: Type of model ('bert', 'roberta', etc.)
+
+        Returns:
+            Dictionary of downloaded file paths
+        """
+        model_dir = self.cache_dir / model_name
+        model_dir.mkdir(parents=True, exist_ok=True)
+
+        files_to_download = {
+            'config.json': f"https://huggingface.co/{model_name}/resolve/main/config.json",
+            'pytorch_model.bin': f"https://huggingface.co/{model_name}/resolve/main/pytorch_model.bin"
+        }
+
+        # Add tokenizer files based on model type
+        if model_type in ['bert', 'roberta', 'distilbert', 'albert']:
+            files_to_download['vocab.txt'] = f"https://huggingface.co/{model_name}/resolve/main/vocab.txt"
+
+        downloaded_files = {}
+
+        for filename, url in files_to_download.items():
+            filepath = model_dir / filename
+            if not filepath.exists():
+                print(f"Downloading {filename}...")
+                try:
+                    self._download_file(url, filepath)
+                except:
+                    
+                    print(f"Error on Downloading A file:\n {filename}\nplease downlaod it manualy and store it as following path:\n{filepath} ")
+                    input("press ENTER when completed .....")
+                    continue 
+                    
+                    
+            else:
+                print(f"Using cached {filename}")
+            downloaded_files[filename] = filepath
+
+        return downloaded_files
+
+    def _download_file(self, url: str, filepath: Path):
+        """Download file with progress bar."""
+        try:
+            response = requests.get(url, stream=True, timeout=30)
+            response.raise_for_status()
+
+            total_size = int(response.headers.get('content-length', 0))
+
+            with open(filepath, 'wb') as f:
+                with tqdm(total=total_size, unit='B', unit_scale=True, desc=filepath.name) as pbar:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        if chunk:
+                            f.write(chunk)
+                            pbar.update(len(chunk))
+        except Exception as e:
+            print(f"Error downloading {url}: {e}")
+            raise
+
+    def load_bert(self, model_name: str, 
+                  config: Optional[ModelConfig] = None,
+                  config_override: Optional[Dict] = None) -> DynamicTransformer:
+        """
+        Load BERT model as shared encoder.
+        
+        Args:
+            model_name: BERT model name (e.g., 'prajjwal1/bert-tiny', 'bert-base-uncased')
+            config: ModelConfig instance
+            config_override: Override configuration parameters
+
+        Returns:
+            DynamicTransformer model with BERT encoder weights
+        """
+        print(f"\n{'='*60}")
+        print(f"Loading BERT model: {model_name}")
+        print(f"{'='*60}")
+        if config and not config.tokenizer_name:
+            config.tokenizer_name=model_name
+        # Download model files
+        files = self.download_model_files(model_name, 'bert')
+
+        # Load config
+        with open(files['config.json'], 'r') as f:
+            bert_config = json.load(f)
+
+        print(f"BERT config:")
+        print(f"  Hidden size: {bert_config.get('hidden_size', 'N/A')}")
+        print(f"  Layers: {bert_config.get('num_hidden_layers', 'N/A')}")
+        print(f"  Attention heads: {bert_config.get('num_attention_heads', 'N/A')}")
+        print(f"  Max position embeddings: {bert_config.get('max_position_embeddings', 'N/A')}")
+        print(f"  Vocabulary size: {bert_config.get('vocab_size', 'N/A')}")
+
+        # Create or update model config
+        if config is None:
+            config = ModelConfig(
+                embed_dim=bert_config.get('hidden_size', 768),
+                num_layers=bert_config.get('num_hidden_layers', 12),
+                num_heads=bert_config.get('num_attention_heads', 12),
+                head_dim=bert_config.get('hidden_size', 768) // bert_config.get('num_attention_heads', 12),
+                ff_mult=4,
+                tokenizer_name=model_name,
+                #use_simple_tokenizer=False,
+                max_seq_len=bert_config.get('max_position_embeddings', 512),
+                dropout=bert_config.get('hidden_dropout_prob', 0.1),
+                #use_rotary_embedding=False,
+                #use_task_adapters=False,
+                special_tokens={}
+            )
+        else:
+            # Update config with BERT settings while preserving user settings
+            #config.tokenizer_name = model_name
+            #config.use_simple_tokenizer = False
+            config.special_tokens = {}
+            
+            #if not hasattr(config, 'embed_dim') or config.embed_dim == 256:
+            config.embed_dim = bert_config.get('hidden_size', config.embed_dim )
+            #if not hasattr(config, 'num_layers') or config.num_layers == 6:
+            config.num_layers = bert_config.get('num_hidden_layers', config.num_layers)
+            #if not hasattr(config, 'num_heads') or config.num_heads == 8:
+            config.num_heads = bert_config.get('num_attention_heads', config.num_heads )
+            config.head_dim = config.embed_dim // config.num_heads
+            #if not hasattr(config, 'max_seq_len') or config.max_seq_len == 256:
+            config.max_seq_len = bert_config.get('max_position_embeddings', config.max_seq_len)
+            config.use_task_adapters=False
+            config.use_rotary_embedding=False
+            
+
+        # Override config if provided
+        if config_override:
+            for key, value in config_override.items():
+                setattr(config, key, value)
+                print(f"Overriding {key}: {value}")
+
+        # Initialize DynamicTransformer
+        print("\nInitializing DynamicTransformer...")
+        model = DynamicTransformer(config)
+        
+        print(f"Tokenizer loaded with vocab size: {len(model.tokenizer)}")
+
+        # Load weights
+        print("\nLoading model weights...")
+        state_dict = torch.load(files['pytorch_model.bin'], map_location='cpu')
+
+        # Map BERT weights to encoder
+        mapped_state_dict = self._map_bert_to_encoder(state_dict, config)
+        
+        # Load mapped weights into encoder only
+        missing_keys, unexpected_keys = model.encoder.load_state_dict(mapped_state_dict, strict=False)
+        
+        if missing_keys:
+            print(f"  Missing keys: {missing_keys}..." if len(missing_keys) > 5 else f"  Missing keys: {missing_keys}")
+        if unexpected_keys:
+            print(f"  Unexpected keys: {unexpected_keys[:5]}..." if len(unexpected_keys) > 5 else f"  Unexpected keys: {unexpected_keys}")
+
+        print(f"\n✓ Successfully loaded BERT model as encoder")
+        print(f"  Encoder parameters: {sum(p.numel() for p in model.encoder.parameters()):,}")
+        print(f"  Total model parameters: {sum(p.numel() for p in model.parameters()):,}")
+        print(f"  Embed dim: {config.embed_dim}")
+        print(f"  Layers: {config.num_layers}")
+        print(f"  Heads: {config.num_heads}")
+        print(f"  Vocabulary size: {len(model.tokenizer)}")
+        print("\n📝 Note: This model has no Tasks: Add tasks using model.add_task() or Train the model on different Tasks to be added")
+
+        return model
+
+    def _map_bert_to_encoder(self, bert_state: Dict, config: ModelConfig) -> Dict:
+        """
+        Map BERT weights to encoder only (no decoders).
+        """
+        mapped = {}
+
+        # Map word embeddings
+        if 'bert.embeddings.word_embeddings.weight' in bert_state:
+            src_weight = bert_state['bert.embeddings.word_embeddings.weight']
+            if src_weight.shape[1] == config.embed_dim:
+                mapped['embedding.weight'] = src_weight
+                print(f"  Mapped word embeddings: shape {src_weight.shape}")
+            else:
+                print(f"  Warning: Word embedding dimension mismatch. Expected {config.embed_dim}, got {src_weight.shape[1]}")
+
+        # Map position embeddings (if needed)
+        if 'bert.embeddings.position_embeddings.weight' in bert_state:
+            # BERT uses learned position embeddings, we use rotary so skip
+            pass
+
+        # Map token type embeddings (if needed)
+        if 'bert.embeddings.token_type_embeddings.weight' in bert_state:
+            # Not used in our model, skip
+            # future work
+            pass
+
+        # Map layer norm before encoder
+        if 'bert.embeddings.LayerNorm.weight' in bert_state:
+            # Not used in our encoder structure, skip
+            # future work
+            pass
+
+        # Map encoder layers
+        for i in range(config.num_layers):  
+            prefix = f'bert.encoder.layer.{i}'
+            target_prefix = f'layers.{i}'
+            
+            # Check if this layer exists in checkpoint
+            if f'{prefix}.attention.self.query.weight' not in bert_state:
+                if i == 0:
+                    print(f"  Warning: Layer {i} not found in checkpoint")
+                continue
+
+            # Self-attention Q, K, V projections
+            if f'{prefix}.attention.self.query.weight' in bert_state:
+                mapped[f'{target_prefix}.attention.q_proj.weight'] = bert_state[f'{prefix}.attention.self.query.weight']
+                mapped[f'{target_prefix}.attention.q_proj.bias'] = bert_state[f'{prefix}.attention.self.query.bias']
+
+                mapped[f'{target_prefix}.attention.k_proj.weight'] = bert_state[f'{prefix}.attention.self.key.weight']
+                mapped[f'{target_prefix}.attention.k_proj.bias'] = bert_state[f'{prefix}.attention.self.key.bias']
+
+                mapped[f'{target_prefix}.attention.v_proj.weight'] = bert_state[f'{prefix}.attention.self.value.weight']
+                mapped[f'{target_prefix}.attention.v_proj.bias'] = bert_state[f'{prefix}.attention.self.value.bias']
+
+            # Attention output projection
+            if f'{prefix}.attention.output.dense.weight' in bert_state:
+                mapped[f'{target_prefix}.attention.out_proj.weight'] = bert_state[f'{prefix}.attention.output.dense.weight']
+                mapped[f'{target_prefix}.attention.out_proj.bias'] = bert_state[f'{prefix}.attention.output.dense.bias']
+
+            # Attention layer norm
+            if f'{prefix}.attention.output.LayerNorm.weight' in bert_state:
+                mapped[f'{target_prefix}.attention_norm.weight'] = bert_state[f'{prefix}.attention.output.LayerNorm.weight']
+                mapped[f'{target_prefix}.attention_norm.bias'] = bert_state[f'{prefix}.attention.output.LayerNorm.bias']
+
+            # Feed-forward intermediate (gate and up projections)
+            if f'{prefix}.intermediate.dense.weight' in bert_state:
+                mapped[f'{target_prefix}.ffn.gate_proj.weight'] = bert_state[f'{prefix}.intermediate.dense.weight']
+                mapped[f'{target_prefix}.ffn.gate_proj.bias'] = bert_state[f'{prefix}.intermediate.dense.bias']
+                mapped[f'{target_prefix}.ffn.up_proj.weight'] = bert_state[f'{prefix}.intermediate.dense.weight']
+                mapped[f'{target_prefix}.ffn.up_proj.bias'] = bert_state[f'{prefix}.intermediate.dense.bias']
+
+            # Feed-forward output
+            if f'{prefix}.output.dense.weight' in bert_state:
+                mapped[f'{target_prefix}.ffn.down_proj.weight'] = bert_state[f'{prefix}.output.dense.weight']
+                mapped[f'{target_prefix}.ffn.down_proj.bias'] = bert_state[f'{prefix}.output.dense.bias']
+
+            # Feed-forward layer norm
+            if f'{prefix}.output.LayerNorm.weight' in bert_state:
+                mapped[f'{target_prefix}.ffn_norm.weight'] = bert_state[f'{prefix}.output.LayerNorm.weight']
+                mapped[f'{target_prefix}.ffn_norm.bias'] = bert_state[f'{prefix}.output.LayerNorm.bias']
+
+            if i < 4:  # Print first few layers
+                print(f"  Mapped layer {i}")
+
+        # Map final layer norm
+        if 'bert.encoder.LayerNorm.weight' in bert_state:
+            mapped['norm.weight'] = bert_state['bert.encoder.LayerNorm.weight']
+            mapped['norm.bias'] = bert_state['bert.encoder.LayerNorm.bias']
+            print("  Mapped final layer norm")
+        else:
+            
+            mapped['norm.weight'] = torch.ones(config.embed_dim)
+            mapped['norm.bias'] = torch.zeros(config.embed_dim)
+            print("  Initialized final layer norm with identity")
+        if 'bert.pooler.dense.weight' in bert_state:
+            # Some BERT variants use pooler instead
+            # future work
+            pass
+
+        return mapped
+
+    def load_roberta(self, model_name: str,
+                     config: Optional[ModelConfig] = None,
+                     config_override: Optional[Dict] = None) -> DynamicTransformer:
+        """
+        Load RoBERTa model as shared encoder.
+        RoBERTa uses the same architecture as BERT with different training.
+        """
+        print(f"\n{'='*60}")
+        print(f"Loading RoBERTa model: {model_name}")
+        print(f"{'='*60}")
+        if config and not config.tokenizer_name:
+            config.tokenizer_name=model_name
+        # Download model files
+        files = self.download_model_files(model_name, 'roberta')
+
+        # Load config
+        with open(files['config.json'], 'r') as f:
+            roberta_config = json.load(f)
+
+        print(f"RoBERTa config:")
+        print(f"  Hidden size: {roberta_config.get('hidden_size', 'N/A')}")
+        print(f"  Layers: {roberta_config.get('num_hidden_layers', 'N/A')}")
+        print(f"  Attention heads: {roberta_config.get('num_attention_heads', 'N/A')}")
+
+        # Create or update model config
+        if config is None:
+            config = ModelConfig(
+                embed_dim=roberta_config.get('hidden_size', 768),
+                num_layers=roberta_config.get('num_hidden_layers', 12),
+                num_heads=roberta_config.get('num_attention_heads', 12),
+                head_dim=roberta_config.get('hidden_size', 768) // roberta_config.get('num_attention_heads', 12),
+                ff_mult=4,
+                tokenizer_name=model_name,
+                #use_simple_tokenizer=False,
+                max_seq_len=roberta_config.get('max_position_embeddings', 514),
+                dropout=roberta_config.get('hidden_dropout_prob', 0.1),
+                use_rotary_embedding=False,
+                use_task_adapters=False,
+                special_tokens={}
+            )
+        else:
+            #config.tokenizer_name = model_name
+            #config.use_simple_tokenizer = False
+            config.special_tokens = {}
+            #if not hasattr(config, 'embed_dim') or config.embed_dim == 256:
+            config.embed_dim = roberta_config.get('hidden_size', config.embed_dim)
+            #if not hasattr(config, 'num_layers') or config.num_layers == 6:
+            config.num_layers = roberta_config.get('num_hidden_layers', config.num_layers)
+            #if not hasattr(config, 'num_heads') or config.num_heads == 8:
+            config.num_heads = roberta_config.get('num_attention_heads', config.num_heads )
+            config.head_dim = config.embed_dim // config.num_heads
+            config.use_task_adapters=False
+            config.use_rotary_embedding=False
+
+        # Override config if provided
+        if config_override:
+            for key, value in config_override.items():
+                setattr(config, key, value)
+                print(f"Overriding {key}: {value}")
+
+        # Initialize DynamicTransformer
+        print("\nInitializing DynamicTransformer...")
+        model = DynamicTransformer(config)
+        
+        print(f"Tokenizer loaded with vocab size: {len(model.tokenizer)}")
+
+        # Load weights
+        print("\nLoading model weights...")
+        state_dict = torch.load(files['pytorch_model.bin'], map_location='cpu')
+
+        # RoBERTa uses similar structure to BERT
+        # Need to map 'roberta.' prefix to 'bert.' prefix
+        renamed_state_dict = {}
+        for key, value in state_dict.items():
+            if key.startswith('roberta.'):
+                new_key = 'bert.' + key[8:]
+                renamed_state_dict[new_key] = value
+            else:
+                renamed_state_dict[key] = value
+
+        mapped_state_dict = self._map_bert_to_encoder(renamed_state_dict, config)
+        model.encoder.load_state_dict(mapped_state_dict, strict=False)
+
+        print(f"\n✓ Successfully loaded RoBERTa model as encoder")
+        print(f"  Encoder parameters: {sum(p.numel() for p in model.encoder.parameters()):,}")
+        print(f"  Embed dim: {config.embed_dim}")
+        print(f"  Layers: {config.num_layers}")
+        print(f"  Vocabulary size: {len(model.tokenizer)}")
+
+        return model
+
+    def load_distilbert(self, model_name: str,
+                        config: Optional[ModelConfig] = None,
+                        config_override: Optional[Dict] = None) -> DynamicTransformer:
+        """
+        Load DistilBERT model as shared encoder.
+        DistilBERT is a distilled version of BERT with fewer layers.
+        """
+        print(f"\n{'='*60}")
+        print(f"Loading DistilBERT model: {model_name}")
+        print(f"{'='*60}")
+
+        # Download model files
+        files = self.download_model_files(model_name, 'distilbert')
+
+        # Load config
+        with open(files['config.json'], 'r') as f:
+            distilbert_config = json.load(f)
+
+        print(f"DistilBERT config:")
+        print(f"  Hidden size: {distilbert_config.get('dim', 'N/A')}")
+        print(f"  Layers: {distilbert_config.get('n_layers', 'N/A')}")
+        print(f"  Attention heads: {distilbert_config.get('n_heads', 'N/A')}")
+
+        # Create or update model config
+        if config is None:
+            config = ModelConfig(
+                embed_dim=distilbert_config.get('dim', 768),
+                num_layers=distilbert_config.get('n_layers', 6),
+                num_heads=distilbert_config.get('n_heads', 12),
+                head_dim=distilbert_config.get('dim', 768) // distilbert_config.get('n_heads', 12),
+                ff_mult=4,
+                tokenizer_name=model_name,
+                #use_simple_tokenizer=False,
+                max_seq_len=distilbert_config.get('max_position_embeddings', 512),
+                dropout=distilbert_config.get('dropout', 0.1),
+                use_rotary_embedding=False,
+                use_task_adapters=False,
+                special_tokens={}
+            )
+        else:
+            #config.tokenizer_name = model_name
+            #config.use_simple_tokenizer = False
+            config.special_tokens = {}
+            #if not hasattr(config, 'embed_dim') or config.embed_dim == 256:
+            config.embed_dim = distilbert_config.get('dim', config.embed_dim)
+            #if not hasattr(config, 'num_layers') or config.num_layers == 6:
+            config.num_layers = distilbert_config.get('n_layers', config.num_layers )
+            #if not hasattr(config, 'num_heads') or config.num_heads == 8:
+            config.num_heads = distilbert_config.get('n_heads', config.num_heads)
+            config.head_dim = config.embed_dim // config.num_heads
+            config.use_task_adapters=False
+            config.use_rotary_embedding=False
+
+        # Override config if provided
+        if config_override:
+            for key, value in config_override.items():
+                setattr(config, key, value)
+                print(f"Overriding {key}: {value}")
+
+        # Initialize DynamicTransformer
+        print("\nInitializing DynamicTransformer...")
+        model = DynamicTransformer(config)
+        
+        print(f"Tokenizer loaded with vocab size: {len(model.tokenizer)}")
+
+        # Load weights
+        print("\nLoading model weights...")
+        state_dict = torch.load(files['pytorch_model.bin'], map_location='cpu')
+
+        # Map DistilBERT weights to encoder
+        mapped_state_dict = self._map_distilbert_to_encoder(state_dict, config)
+        
+        # Load mapped weights
+        model.encoder.load_state_dict(mapped_state_dict, strict=False)
+
+        print(f"\n✓ Successfully loaded DistilBERT model as encoder")
+        print(f"  Encoder parameters: {sum(p.numel() for p in model.encoder.parameters()):,}")
+        print(f"  Embed dim: {config.embed_dim}")
+        print(f"  Layers: {config.num_layers}")
+        print(f"  Vocabulary size: {len(model.tokenizer)}")
+
+        return model
+
+    def _map_distilbert_to_encoder(self, state_dict: Dict, config: ModelConfig) -> Dict:
+        """
+        Map DistilBERT weights to encoder.
+        DistilBERT has a simpler structure than BERT.
+        """
+        mapped = {}
+
+        # Word embeddings
+        if 'distilbert.embeddings.word_embeddings.weight' in state_dict:
+            src_weight = state_dict['distilbert.embeddings.word_embeddings.weight']
+            if src_weight.shape[1] == config.embed_dim:
+                mapped['embedding.weight'] = src_weight
+                print(f"  Mapped word embeddings: shape {src_weight.shape}")
+
+        # Position embeddings (if present)
+        if 'distilbert.embeddings.position_embeddings.weight' in state_dict:
+            # DistilBERT uses learned position embeddings, we skip
+            # future work
+            pass
+
+        # Transformer layers
+        for i in range(config.num_layers):
+            prefix = f'distilbert.transformer.layer.{i}'
+            target_prefix = f'layers.{i}'
+            
+            # Check if layer exists
+            if f'{prefix}.attention.q_lin.weight' not in state_dict:
+                continue
+
+            # Q, K, V projections (DistilBERT has separate linear layers)
+            if f'{prefix}.attention.q_lin.weight' in state_dict:
+                mapped[f'{target_prefix}.attention.q_proj.weight'] = state_dict[f'{prefix}.attention.q_lin.weight']
+                mapped[f'{target_prefix}.attention.q_proj.bias'] = state_dict[f'{prefix}.attention.q_lin.bias']
+                
+                mapped[f'{target_prefix}.attention.k_proj.weight'] = state_dict[f'{prefix}.attention.k_lin.weight']
+                mapped[f'{target_prefix}.attention.k_proj.bias'] = state_dict[f'{prefix}.attention.k_lin.bias']
+                
+                mapped[f'{target_prefix}.attention.v_proj.weight'] = state_dict[f'{prefix}.attention.v_lin.weight']
+                mapped[f'{target_prefix}.attention.v_proj.bias'] = state_dict[f'{prefix}.attention.v_lin.bias']
+
+            # Attention output projection
+            if f'{prefix}.attention.out_lin.weight' in state_dict:
+                mapped[f'{target_prefix}.attention.out_proj.weight'] = state_dict[f'{prefix}.attention.out_lin.weight']
+                mapped[f'{target_prefix}.attention.out_proj.bias'] = state_dict[f'{prefix}.attention.out_lin.bias']
+
+            # Self-attention layer norm
+            if f'{prefix}.sa_layer_norm.weight' in state_dict:
+                mapped[f'{target_prefix}.attention_norm.weight'] = state_dict[f'{prefix}.sa_layer_norm.weight']
+                mapped[f'{target_prefix}.attention_norm.bias'] = state_dict[f'{prefix}.sa_layer_norm.bias']
+
+            # Feed-forward first linear (gate and up share same weights)
+            if f'{prefix}.ffn.lin1.weight' in state_dict:
+                mapped[f'{target_prefix}.ffn.gate_proj.weight'] = state_dict[f'{prefix}.ffn.lin1.weight']
+                mapped[f'{target_prefix}.ffn.gate_proj.bias'] = state_dict[f'{prefix}.ffn.lin1.bias']
+                mapped[f'{target_prefix}.ffn.up_proj.weight'] = state_dict[f'{prefix}.ffn.lin1.weight']
+                mapped[f'{target_prefix}.ffn.up_proj.bias'] = state_dict[f'{prefix}.ffn.lin1.bias']
+
+            # Feed-forward second linear
+            if f'{prefix}.ffn.lin2.weight' in state_dict:
+                mapped[f'{target_prefix}.ffn.down_proj.weight'] = state_dict[f'{prefix}.ffn.lin2.weight']
+                mapped[f'{target_prefix}.ffn.down_proj.bias'] = state_dict[f'{prefix}.ffn.lin2.bias']
+
+            # Feed-forward layer norm
+            if f'{prefix}.output_layer_norm.weight' in state_dict:
+                mapped[f'{target_prefix}.ffn_norm.weight'] = state_dict[f'{prefix}.output_layer_norm.weight']
+                mapped[f'{target_prefix}.ffn_norm.bias'] = state_dict[f'{prefix}.output_layer_norm.bias']
+
+            if i < 2:
+                print(f"  Mapped layer {i}")
+
+        # Final layer norm (if exists)
+        if 'distilbert.transformer.LayerNorm.weight' in state_dict:
+            mapped['norm.weight'] = state_dict['distilbert.transformer.LayerNorm.weight']
+            mapped['norm.bias'] = state_dict['distilbert.transformer.LayerNorm.bias']
+            print(f"  Mapped final layer norm")
+
+        return mapped
+
+    def load_albert(self, model_name: str,
+                    config: Optional[ModelConfig] = None,
+                    config_override: Optional[Dict] = None) -> DynamicTransformer:
+        """
+        Load ALBERT model as shared encoder.
+        ALBERT uses parameter sharing across layers.
+        """
+        print(f"\n{'='*60}")
+        print(f"Loading ALBERT model: {model_name}")
+        print(f"{'='*60}")
+
+        # Download model files
+        files = self.download_model_files(model_name, 'albert')
+
+        # Load config
+        with open(files['config.json'], 'r') as f:
+            albert_config = json.load(f)
+
+        print(f"ALBERT config:")
+        print(f"  Hidden size: {albert_config.get('hidden_size', 'N/A')}")
+        print(f"  Layers: {albert_config.get('num_hidden_layers', 'N/A')}")
+        print(f"  Attention heads: {albert_config.get('num_attention_heads', 'N/A')}")
+        print(f"  Embedding size: {albert_config.get('embedding_size', 'N/A')}")
+
+        # Create or update model config
+        if config is None:
+            config = ModelConfig(
+                embed_dim=albert_config.get('hidden_size', 768),
+                num_layers=albert_config.get('num_hidden_layers', 12),
+                num_heads=albert_config.get('num_attention_heads', 12),
+                head_dim=albert_config.get('hidden_size', 768) // albert_config.get('num_attention_heads', 12),
+                ff_mult=4,
+                tokenizer_name=model_name,
+                #use_simple_tokenizer=False,
+                max_seq_len=albert_config.get('max_position_embeddings', 512),
+                dropout=albert_config.get('hidden_dropout_prob', 0.1),
+                use_rotary_embedding=False,
+                use_task_adapters=False,
+                special_tokens={}
+            )
+        else:
+            #config.tokenizer_name = model_name
+            #config.use_simple_tokenizer = False
+            config.special_tokens = {}
+            
+            config.embed_dim = albert_config.get('hidden_size', config.embed_dim)
+            
+            config.num_layers = albert_config.get('num_hidden_layers', config.num_layers)
+            
+            config.num_heads = albert_config.get('num_attention_heads', config.num_heads)
+            config.head_dim = config.embed_dim // config.num_heads
+            config.use_task_adapters=False
+            config.use_rotary_embedding=False
+            
+
+        # Override config if provided
+        if config_override:
+            for key, value in config_override.items():
+                setattr(config, key, value)
+                print(f"Overriding {key}: {value}")
+
+        # Initialize DynamicTransformer
+        print("\nInitializing DynamicTransformer...")
+        model = DynamicTransformer(config)
+        
+        print(f"Tokenizer loaded with vocab size: {len(model.tokenizer)}")
+
+        # Load weights
+        print("\nLoading model weights...")
+        state_dict = torch.load(files['pytorch_model.bin'], map_location='cpu')
+
+        # ALBERT has a similar structure to BERT but with embedding projection
+        # We'll map using BERT mapping with prefix adjustments
+        renamed_state_dict = {}
+        for key, value in state_dict.items():
+            if key.startswith('albert.'):
+                new_key = 'bert.' + key[7:]
+                renamed_state_dict[new_key] = value
+            else:
+                renamed_state_dict[key] = value
+
+        mapped_state_dict = self._map_bert_to_encoder(renamed_state_dict, config)
+        model.encoder.load_state_dict(mapped_state_dict, strict=False)
+
+        print(f"\n✓ Successfully loaded ALBERT model as encoder")
+        print(f"  Encoder parameters: {sum(p.numel() for p in model.encoder.parameters()):,}")
+        print(f"  Embed dim: {config.embed_dim}")
+        print(f"  Layers: {config.num_layers} (note: ALBERT shares parameters across layers)")
+        print(f"  Vocabulary size: {len(model.tokenizer)}")
+
+        return model
+
+    def list_available_models(self) -> Dict[str, List[str]]:
+        """
+        List available pretrained models for encoder loading.
+        
+        Returns:
+            Dictionary of model types with list of recommended models
+        """
+        return {
+            'bert': [
+                'prajjwal1/bert-tiny',      # 2 layers, 128 dim (~4.4M params)
+                'prajjwal1/bert-mini',      # 4 layers, 256 dim (~11M params)
+                'prajjwal1/bert-small',     # 4 layers, 512 dim (~28M params)
+                'prajjwal1/bert-medium',    # 8 layers, 512 dim (~41M params)
+                'bert-base-uncased',        # 12 layers, 768 dim (~110M params)
+                'bert-large-uncased',       # 24 layers, 1024 dim (~336M params)
+                'bert-base-multilingual-cased',  # Multilingual support
+                'bert-base-arabic'          # Arabic BERT
+            ],
+            'roberta': [
+                'roberta-base',             # 12 layers, 768 dim (~125M params)
+                'roberta-large',            # 24 layers, 1024 dim (~355M params)
+                #'xlm-roberta-base',         # Multilingual RoBERTa
+                #'xlm-roberta-large'         # Large multilingual
+            ],
+            'distilbert': [
+                'distilbert-base-uncased',  # 6 layers, 768 dim (~66M params)
+                'distilbert-base-multilingual-cased'  # Multilingual
+            ],
+            'albert': [
+                'albert-base-v2',           # 12 layers, 768 dim (~12M params)
+                'albert-large-v2',          # 24 layers, 1024 dim (~18M params)
+                'albert-xlarge-v2',         # 24 layers, 2048 dim (~60M params)
+                'albert-xxlarge-v2'         # 12 layers, 4096 dim (~235M params)
+            ]
+        }
+
+    def get_model_info(self, model_name: str) -> Dict:
+        """
+        Get information about a pretrained model without loading it.
+        
+        Args:
+            model_name: HuggingFace model name
+            
+        Returns:
+            Dictionary with model information
+        """
+        try:
+            model_type = self._detect_model_type(model_name)
+            
+            # Try to download config
+            model_dir = self.cache_dir / model_name
+            config_path = model_dir / 'config.json'
+            
+            if not config_path.exists():
+                url = f"https://huggingface.co/{model_name}/resolve/main/config.json"
+                response = requests.get(url, timeout=10)
+                if response.status_code == 200:
+                    config = response.json()
+                else:
+                    return {"error": "Could not fetch config"}
+            else:
+                with open(config_path, 'r') as f:
+                    config = json.load(f)
+            
+            info = {
+                "model_name": model_name,
+                "model_type": model_type,
+                "supported": model_type in self.SUPPORTED_MODELS,
+                "architecture": {}
+            }
+            
+            if model_type == 'bert':
+                info["architecture"] = {
+                    "hidden_size": config.get('hidden_size', 'N/A'),
+                    "num_layers": config.get('num_hidden_layers', 'N/A'),
+                    "num_heads": config.get('num_attention_heads', 'N/A'),
+                    "vocab_size": config.get('vocab_size', 'N/A'),
+                    "max_position_embeddings": config.get('max_position_embeddings', 'N/A')
+                }
+            elif model_type == 'roberta':
+                info["architecture"] = {
+                    "hidden_size": config.get('hidden_size', 'N/A'),
+                    "num_layers": config.get('num_hidden_layers', 'N/A'),
+                    "num_heads": config.get('num_attention_heads', 'N/A'),
+                    "vocab_size": config.get('vocab_size', 'N/A'),
+                    "max_position_embeddings": config.get('max_position_embeddings', 'N/A')
+                }
+            elif model_type == 'distilbert':
+                info["architecture"] = {
+                    "hidden_size": config.get('dim', 'N/A'),
+                    "num_layers": config.get('n_layers', 'N/A'),
+                    "num_heads": config.get('n_heads', 'N/A'),
+                    "vocab_size": config.get('vocab_size', 'N/A'),
+                    "max_position_embeddings": config.get('max_position_embeddings', 'N/A')
+                }
+            elif model_type == 'albert':
+                info["architecture"] = {
+                    "hidden_size": config.get('hidden_size', 'N/A'),
+                    "num_layers": config.get('num_hidden_layers', 'N/A'),
+                    "num_heads": config.get('num_attention_heads', 'N/A'),
+                    "vocab_size": config.get('vocab_size', 'N/A'),
+                    "embedding_size": config.get('embedding_size', 'N/A')
+                }
+            
+            return info
+            
+        except Exception as e:
+            return {"error": str(e)}

src/dytr/tokenization/__init__.py

@@ -0,0 +1,41 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Tokenization module for Dynamic Transformers.
+
+This module provides tokenizer implementations and download utilities.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+from dytr.tokenization.download_manager import (
+    DownloadManager,
+    download_file,
+    download_tokenizer_files,
+    get_url_from_HF,
+    process_vocab_text,
+)
+from dytr.tokenization.simple_tokenizer import SimpleTokenizer
+
+__all__ = [
+    "SimpleTokenizer",
+    "DownloadManager",
+    "download_file",
+    "download_tokenizer_files",
+    "get_url_from_HF",
+    "process_vocab_text",
+]

src/dytr/tokenization/download_manager.py

@@ -0,0 +1,280 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Download manager for tokenizer files from HuggingFace.
+
+This module provides functionality to download tokenizer files from HuggingFace
+with resume support and caching.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+import hashlib
+import json
+import pickle
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+import pandas as pd
+import requests
+from tqdm import tqdm
+
+
+def get_url_from_HF(HF_path: str):
+    """
+    Generate download URLs for HuggingFace tokenizer files.
+
+    Args:
+        HF_path: HuggingFace model path (e.g., 'bert-base-uncased')
+
+    Returns:
+        Tuple of (directory_name, vocab_url, special_tokens_map_url)
+    """
+    vocab_url = f"https://huggingface.co/{HF_path}/resolve/main/vocab.txt?download=true"
+    special_token_map = (
+        f"https://huggingface.co/{HF_path}/resolve/main/special_tokens_map.json?download=true"
+    )
+    file_dir = HF_path.split("/")[-1].replace("-", "_")
+    return file_dir, vocab_url, special_token_map
+
+
+def process_vocab_text(path: Union[str, Path]) -> Dict[str, int]:
+    """
+    Process vocabulary text file into token-to-id mapping.
+
+    Args:
+        path: Path to vocabulary text file
+
+    Returns:
+        Dictionary mapping tokens to IDs
+    """
+    path = Path(path)
+    json_path = path.with_suffix(".json")
+
+    if json_path.exists():
+        with open(json_path, "r", encoding="utf-8") as f:
+            token_to_id = json.load(f)
+        print("JSON OF VOCAB ALREADY SAVED")
+        return token_to_id
+
+    print("*" * 60)
+    with open(path, "r", encoding="utf-8") as f:
+        data = f.read()
+
+    token_to_id = {k: idx for idx, k in enumerate(data.split("\n")) if k}
+    return token_to_id
+
+
+class DownloadManager:
+    """
+    Manages file downloads with resume support and caching.
+
+    This class handles downloading files from URLs, supports resuming
+    interrupted downloads, and caches downloaded files to avoid
+    re-downloading.
+
+    Args:
+        base_dir: Base directory for storing downloads
+        registry_file: Name of the registry file for tracking downloads
+    """
+
+    def __init__(self, base_dir: str = "downloads", registry_file: str = "download_registry.pkl"):
+        self.base_dir = Path(base_dir)
+        self.base_dir.mkdir(exist_ok=True)
+        self.registry_file = self.base_dir / registry_file
+        self.registry = self._load_registry()
+
+    def _load_registry(self) -> Dict[str, str]:
+        """Load download registry from disk."""
+        if self.registry_file.exists():
+            with open(self.registry_file, "rb") as f:
+                return pickle.load(f)
+        return {}
+
+    def _save_registry(self):
+        """Save download registry to disk."""
+        with open(self.registry_file, "wb") as f:
+            pickle.dump(self.registry, f)
+
+    def _get_file_path(self, url: str, subfolder: Optional[str] = None) -> Path:
+        """
+        Generate file path for a given URL.
+
+        Args:
+            url: Download URL
+            subfolder: Optional subfolder within base directory
+
+        Returns:
+            Path object for the file
+        """
+        url_hash = hashlib.md5(url.encode()).hexdigest()[:8]
+        original_filename = url.split("/")[-1].split("?")[0]
+
+        if not original_filename:
+            original_filename = f"download_{url_hash}"
+
+        if subfolder:
+            file_path = self.base_dir / subfolder / original_filename
+            file_path.parent.mkdir(parents=True, exist_ok=True)
+        else:
+            file_path = self.base_dir / original_filename
+
+        return file_path
+
+    def download_file(
+        self, url: str, subfolder: Optional[str] = None
+    ) -> Union[str, pd.DataFrame, dict, list]:
+        """
+        Download a file from URL with resume support.
+
+        Args:
+            url: URL to download
+            subfolder: Optional subfolder to store the file
+
+        Returns:
+            Parsed file content (dict for JSON, DataFrame for CSV/Excel, str for text)
+        """
+        file_path = self._get_file_path(url, subfolder)
+
+        # Check if file already exists in registry
+        if str(file_path) in self.registry.values():
+            if file_path.exists():
+                print(f"File already exists: {file_path}")
+                return self._read_file(file_path)
+
+        # Check if file exists directly
+        if file_path.exists():
+            print(f"File already exists: {file_path}")
+            return self._read_file(file_path)
+
+        # Download file
+        self._download_with_resume(url, file_path)
+
+        # Update registry
+        self.registry[url] = str(file_path)
+        self._save_registry()
+
+        return self._read_file(file_path)
+
+    def _download_with_resume(self, url: str, file_path: Path):
+        """
+        Download file with resume support for interrupted downloads.
+
+        Args:
+            url: URL to download
+            file_path: Path where to save the file
+        """
+        existing_size = file_path.stat().st_size if file_path.exists() else 0
+
+        headers = {"Range": f"bytes={existing_size}-"} if existing_size > 0 else {}
+
+        response = requests.get(url, headers=headers, stream=True)
+        response.raise_for_status()
+
+        total_size = int(response.headers.get("content-length", 0)) + existing_size
+
+        mode = "ab" if existing_size > 0 else "wb"
+
+        with open(file_path, mode) as f:
+            with tqdm(
+                total=total_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=f"Downloading {file_path.name}",
+                initial=existing_size,
+                bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
+            ) as pbar:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+
+    def _read_file(self, file_path: Path) -> Union[str, pd.DataFrame, dict, list]:
+        """
+        Read and parse a downloaded file based on its extension.
+
+        Args:
+            file_path: Path to the file
+
+        Returns:
+            Parsed file content
+        """
+        file_extension = file_path.suffix.lower()
+
+        if file_extension == ".json":
+            with open(file_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        elif file_extension == ".csv":
+            return pd.read_csv(file_path)
+        elif file_extension == ".txt":
+            if "vocab" in str(file_path).lower():
+                return process_vocab_text(file_path)
+            with open(file_path, "r", encoding="utf-8") as f:
+                return f.read()
+        elif file_extension in [".xlsx", ".xls"]:
+            return pd.read_excel(file_path)
+        else:
+            with open(file_path, "r", encoding="utf-8") as f:
+                return f.read()
+
+    def get_downloaded_files(self) -> Dict[str, str]:
+        """
+        Get all downloaded files from registry.
+
+        Returns:
+            Dictionary mapping URLs to file paths
+        """
+        return self.registry.copy()
+
+
+def download_file(
+    url: str, subfolder: Optional[str] = None, base_dir: str = "downloads"
+) -> Union[str, pd.DataFrame, dict, list]:
+    """
+    Convenience function to download a single file.
+
+    Args:
+        url: URL to download
+        subfolder: Optional subfolder for the file
+        base_dir: Base directory for downloads
+
+    Returns:
+        Parsed file content
+    """
+    manager = DownloadManager(base_dir)
+    return manager.download_file(url, subfolder)
+
+
+def download_tokenizer_files(HF_tokenizer_path: str) -> tuple:
+    """
+    Download tokenizer files from HuggingFace.
+
+    Args:
+        HF_tokenizer_path: HuggingFace tokenizer path (e.g., 'bert-base-uncased')
+
+    Returns:
+        Tuple of (vocab_dict, special_tokens_map)
+    """
+    sub_dir, vocab_path, special_token_map_path = get_url_from_HF(HF_tokenizer_path)
+    vocab = download_file(vocab_path, sub_dir)
+    try:
+        special_token_map = download_file(special_token_map_path, sub_dir)
+    except:
+        special_token_map = {}
+
+    return vocab, special_token_map

src/dytr/tokenization/simple_tokenizer.py

@@ -0,0 +1,839 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Simple tokenizer implementation with multiple tokenization strategies.
+
+This module provides a lightweight tokenizer that supports various tokenization
+methods including WordPiece, BPE, Unigram, and Byte-Level encoding.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+import json
+import os
+from typing import Any, Dict, List, Optional, Union
+
+
+class SimpleTokenizer:
+    """
+    A flexible tokenizer supporting multiple tokenization strategies.
+
+    Supports:
+    - WordPiece tokenization
+    - Byte-Pair Encoding (BPE)
+    - Unigram tokenization
+    - Character-level tokenization
+    - Space-based tokenization
+    - Byte-level tokenization (GPT-2 style)
+
+    Args:
+        vocab: Vocabulary dictionary mapping tokens to IDs
+        special_tokens_map: Dictionary of special tokens
+        tokenizer_type: Type of tokenizer ('wordpiece', 'bpe', 'unigram', 'char', 'space', 'byte_level')
+    """
+
+    def __init__(
+        self,
+        vocab: Dict[str, int],
+        special_tokens_map: Optional[Dict] = None,
+        tokenizer_type: str = "wordpiece", add_tab_newline_vocab: bool =False
+    ):
+        self.vocab = vocab
+        self.inv_vocab = {v: k for k, v in vocab.items()}
+        self.special_tokens_map = special_tokens_map or {}
+        self.tokenizer_type = tokenizer_type
+        self.add_tab_newline_vocab = add_tab_newline_vocab
+
+        # Special tokens
+        self.pad_token = self.special_tokens_map.get("pad_token", "[PAD]")
+        self.unk_token = self.special_tokens_map.get("unk_token", "[UNK]")
+        self.cls_token = self.special_tokens_map.get("cls_token", "[CLS]")
+        self.sep_token = self.special_tokens_map.get("sep_token", "[SEP]")
+        self.mask_token = self.special_tokens_map.get("mask_token", "[MASK]")
+        self.bos_token = self.special_tokens_map.get("bos_token", "<s>")
+        self.eos_token = self.special_tokens_map.get("eos_token", "</s>")
+
+        # Token IDs
+        self.pad_token_id = self.vocab.get(self.pad_token)
+        self.unk_token_id = self.vocab.get(self.unk_token)
+        self.cls_token_id = self.vocab.get(self.cls_token)
+        self.sep_token_id = self.vocab.get(self.sep_token)
+        self.mask_token_id = self.vocab.get(self.mask_token)
+
+        # Add BOS and EOS if not present
+        if self.bos_token not in self.vocab:
+            self.bos_token=self.cls_token
+            
+            #self.add_tokens([self.bos_token])
+        if self.eos_token not in self.vocab:
+            self.eos_token=self.sep_token
+            #self.add_tokens([self.eos_token])
+
+        self.bos_token_id = self.vocab.get(self.bos_token)
+        self.eos_token_id = self.vocab.get(self.eos_token)
+
+        # All special tokens
+        self.all_special_tokens = [
+            self.pad_token,
+            self.unk_token,
+            self.cls_token,
+            self.sep_token,
+            self.mask_token,
+            self.bos_token,
+            self.eos_token,
+        ]
+        self.all_special_ids = [
+            self.pad_token_id,
+            self.unk_token_id,
+            self.cls_token_id,
+            self.sep_token_id,
+            self.mask_token_id,
+            self.bos_token_id,
+            self.eos_token_id,
+        ]
+        
+        if self.add_tab_newline_vocab:
+            # Add newline and tab representations
+            self.add_tokens(["ĉ", "Ċ","Ġ"])
+        self.name_or_path = "local_tokenizer"
+
+        # Tokenization function mapping
+        self.tokenization_functions = {
+            "wordpiece": self._wordpiece_tokenize,
+            "bpe": self._bpe_tokenize,
+            "unigram": self._unigram_tokenize,
+            "char": self._char_tokenize,
+            "space": self._space_tokenize,
+            "byte_level": self._byte_level_tokenize,
+        }
+
+        # BPE merges
+        self.merges = self.special_tokens_map.get("merges", [])
+
+        # Byte-level encoder
+        self.byte_encoder = self._build_byte_encoder() if tokenizer_type == "byte_level" else None
+
+    def _build_byte_encoder(self) -> Dict[str, str]:
+        """Build byte-level encoder for GPT-2 style tokenization."""
+        # Printable ASCII range
+        bs = list(range(ord("!"), ord("~") + 1))
+        bs += list(range(ord("¡"), ord("¬") + 1))
+        bs += list(range(ord("®"), ord("ÿ") + 1))
+        cs = bs[:]
+        n = 0
+        for b in range(2**8):
+            if b not in bs:
+                bs.append(b)
+                cs.append(2**8 + n)
+                n += 1
+        return {chr(b): chr(c) for b, c in zip(bs, cs)}
+
+    def tokenize(self, text: str) -> List[str]:
+        """Tokenize text using the configured tokenization strategy."""
+        if not text:
+            return []
+
+        tokenizer_func = self.tokenization_functions.get(
+            self.tokenizer_type, self._wordpiece_tokenize
+        )
+        return tokenizer_func(text)
+
+    def _space_tokenize(self, text: str) -> List[str]:
+        """Simple space-based tokenization."""
+        return text.split()
+
+    def _char_tokenize(self, text: str) -> List[str]:
+        """Character-level tokenization."""
+        tokens = []
+        for char in text:
+            if char in self.vocab:
+                tokens.append(char)
+            else:
+                tokens.append(self.unk_token)
+        return tokens
+
+    def _wordpiece_tokenize(self, text: str) -> List[str]:
+        """WordPiece tokenization."""
+        if self.add_tab_newline_vocab:
+            text = text.replace("\n", "Ċ").replace("    ", "\t").replace("\t", "ĉ")
+        words = text.split(" ")
+        tokens = []
+
+        for word in words:
+            if word in self.vocab:
+                tokens.append(word)
+                continue
+            subwords = self._wordpiece_split(word)
+            tokens.extend(subwords)
+
+        return tokens
+
+    def _wordpiece_split(self, word: str) -> List[str]:
+        """Split a word into WordPiece subwords."""
+        if not word:
+            return []
+        if word in self.vocab:
+            return [word]
+
+        tokens = []
+        remaining = word
+        start = 0
+
+        while start < len(remaining):
+            end = len(remaining)
+            best_token = None
+
+            while start < end:
+                subword = remaining[start:end]
+                if start > 0:
+                    subword = "##" + subword
+                if subword in self.vocab:
+                    best_token = subword
+                    break
+                end -= 1
+
+            if best_token is None:
+                # Unknown characters
+                for char in remaining[start:]:
+                    if char in self.vocab:
+                        tokens.append(char)
+                    else:
+                        tokens.append(self.unk_token)
+                break
+            else:
+                tokens.append(best_token)
+                start = end
+
+        return tokens
+
+    def _bpe_tokenize(self, text: str) -> List[str]:
+        """Byte-Pair Encoding tokenization."""
+        words = self._preprocess_text(text)
+        tokens = []
+
+        for word in words:
+            if word in self.vocab:
+                tokens.append(word)
+            else:
+                bpe_tokens = self._apply_bpe(word)
+                tokens.extend(bpe_tokens)
+
+        return tokens
+
+    def _preprocess_text(self, text: str) -> List[str]:
+        """Preprocess text for BPE tokenization."""
+        processed = []
+        for i, char in enumerate(text):
+            if char == " ":
+                processed.append("Ġ")
+            elif char == "\n":
+                count = 1
+                idx = i + 1
+                while idx < len(text) and text[idx] == "\n":
+                    count += 1
+                    idx += 1
+                processed.append("Ċ" * count)
+            elif char == "\t":
+                processed.append("ĉ")
+            else:
+                processed.append(char)
+        return processed
+
+    def _apply_bpe(self, word: Union[str, List[str]]) -> List[str]:
+        """Apply BPE merges to a word."""
+        if isinstance(word, list):
+            word = "".join(word)
+
+        if not self.merges:
+            return [word] if word in self.vocab else [self.unk_token]
+
+        tokens = list(word)
+
+        for merge in self.merges:
+            pair = merge.split()
+            if len(pair) != 2:
+                continue
+
+            new_tokens = []
+            i = 0
+            while i < len(tokens):
+                if i < len(tokens) - 1 and tokens[i] == pair[0] and tokens[i + 1] == pair[1]:
+                    new_tokens.append("".join(pair))
+                    i += 2
+                else:
+                    new_tokens.append(tokens[i])
+                    i += 1
+            tokens = new_tokens
+
+        return [token if token in self.vocab else self.unk_token for token in tokens]
+
+    def _unigram_tokenize(self, text: str) -> List[str]:
+        """Unigram tokenization using Viterbi algorithm."""
+        chars = list(text)
+        n = len(chars)
+        dp = [-1] * (n + 1)
+        dp[0] = 0
+
+        # Find best segmentation
+        for i in range(n):
+            if dp[i] == -1:
+                continue
+            for j in range(i + 1, min(i + 20, n + 1)):
+                token = "".join(chars[i:j])
+                if token in self.vocab or token in self.all_special_tokens:
+                    if dp[j] == -1 or dp[j] > dp[i] + 1:
+                        dp[j] = dp[i] + 1
+
+        if dp[n] == -1:
+            return [self.unk_token]
+
+        # Reconstruct tokens
+        tokens = []
+        i = n
+        while i > 0:
+            for j in range(max(0, i - 20), i):
+                if dp[j] != -1 and dp[i] == dp[j] + 1:
+                    token = "".join(chars[j:i])
+                    if token in self.vocab:
+                        tokens.insert(0, token)
+                    else:
+                        tokens.insert(0, self.unk_token)
+                    i = j
+                    break
+
+        return tokens
+
+    def _byte_level_tokenize(self, text: str) -> List[str]:
+        """Byte-level tokenization (GPT-2 style)."""
+        if not self.byte_encoder:
+            return self._char_tokenize(text)
+
+        # Encode bytes
+        bytes_text = text.encode("utf-8")
+        encoded_chars = []
+        for b in bytes_text:
+            if b < 128:
+                encoded_chars.append(chr(b))
+            else:
+                encoded_chars.append(self.byte_encoder.get(chr(b), chr(b)))
+
+        # Preprocess
+        processed = []
+        i = 0
+        while i < len(encoded_chars):
+            if encoded_chars[i] == " ":
+                processed.append("Ġ")
+                i += 1
+            elif encoded_chars[i] == "\n":
+                count = 1
+                while i + count < len(encoded_chars) and encoded_chars[i + count] == "\n":
+                    count += 1
+                processed.append("Ċ" * count)
+                i += count
+            elif encoded_chars[i] == "\t":
+                processed.append("ĉ")
+                i += 1
+            else:
+                processed.append(encoded_chars[i])
+                i += 1
+
+        # Apply BPE if merges exist
+        if self.merges:
+            return self._apply_bpe(processed)
+
+        return [p if p in self.vocab else self.unk_token for p in processed]
+
+    def convert_tokens_to_ids(self, tokens: Union[str, List[str]]) -> List[int]:
+        """Convert tokens to token IDs."""
+        if isinstance(tokens, str):
+            tokens = [tokens]
+        return [self.vocab.get(token, self.unk_token_id) for token in tokens]
+
+    def convert_ids_to_tokens(
+        self, ids: Union[int, List[int]], skip_special_tokens: bool = False
+    ) -> List[str]:
+        """Convert token IDs to tokens."""
+        if hasattr(ids, "cpu"):
+            ids = ids.cpu().tolist()
+        elif isinstance(ids, int):
+            ids = [ids]
+
+        tokens = []
+        for idx in ids:
+            token = self.inv_vocab.get(idx, self.unk_token)
+            if skip_special_tokens and token in self.all_special_tokens:
+                continue
+            tokens.append(token)
+
+        return tokens
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Convert a list of tokens to a string."""
+        text = ""
+        for i, token in enumerate(tokens):
+            if token.startswith("##"):
+                text += token[2:]
+            elif token in ["Ġ", "Ċ", "ĉ"] or (
+                self.tokenizer_type == "byte_level" and token in ["Ġ", "Ċ", "ĉ"]
+            ):
+                if token == "Ġ":
+                    text += " "
+                elif token == "Ċ":
+                    text += "\n"
+                elif token == "ĉ":
+                    text += "\t"
+            elif token in self.all_special_tokens:
+                if token == self.sep_token and i == len(tokens) - 1:
+                    continue
+                text += " " + token
+            else:
+                if (
+                    text
+                    and not text.endswith(" ")
+                    and not text.endswith("\n")
+                    and not text.endswith("\t")
+                ):
+                    text += " "
+                text += token
+
+        text = text.strip()
+        for special in self.all_special_tokens:
+            text = text.replace(" " + special, "")
+            text = text.replace(special + " ", "")
+
+        return text.strip()
+
+    def _get_word_ids(
+        self, input_ids: List[int], offset_mapping: Optional[List] = None
+    ) -> List[int]:
+        """Get word IDs for token classification."""
+        if offset_mapping is not None:
+            word_ids = []
+            current_word_idx = 0
+            last_end = 0
+
+            for i, (start, end) in enumerate(offset_mapping):
+                if start == 0 and end == 0:
+                    word_ids.append(None)
+                else:
+                    if i > 0 and start > last_end:
+                        current_word_idx += 1
+                    word_ids.append(current_word_idx)
+                    last_end = end
+
+            return word_ids
+
+        tokens = self.convert_ids_to_tokens(input_ids)
+        word_ids = []
+        current_word_idx = 0
+
+        for i, token in enumerate(tokens):
+            if token.startswith("##"):
+                word_ids.append(current_word_idx)
+            else:
+                if i > 0 and not tokens[i - 1].startswith("##"):
+                    current_word_idx += 1
+                word_ids.append(current_word_idx)
+
+        return word_ids
+
+    def word_ids(self, batch_encoding: Any) -> Optional[List[int]]:
+        """Extract word IDs from a batch encoding."""
+        if hasattr(batch_encoding, "encodings"):
+            word_ids_list = []
+            for encoding in batch_encoding.encodings:
+                word_ids_list.append(
+                    self._get_word_ids(encoding.ids, getattr(encoding, "offsets", None))
+                )
+            return word_ids_list
+        elif isinstance(batch_encoding, dict) and "input_ids" in batch_encoding:
+            return self._get_word_ids(
+                batch_encoding["input_ids"], batch_encoding.get("offset_mapping")
+            )
+        return None
+
+    def encode(
+        self,
+        text: str,
+        add_special_tokens: bool = False,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        padding: bool = False,
+        return_tensors: Optional[str] = None,
+        return_token_type_ids: bool = False,
+        return_attention_mask: bool = True,
+        return_offsets_mapping: bool = False,
+    ) -> Dict:
+        """Encode text to token IDs."""
+        tokens = self.tokenize(text)
+
+        # Add special tokens
+        if add_special_tokens:
+            if self.bos_token and self.eos_token:
+                tokens = [self.bos_token] + tokens + [self.eos_token]
+            else:
+                tokens = [self.cls_token] + tokens + [self.sep_token]
+
+        input_ids = self.convert_tokens_to_ids(tokens)
+        attention_mask = [1] * len(input_ids)
+
+        # Build offset mapping for token classification
+        offset_mapping = None
+        if return_offsets_mapping:
+            offset_mapping = []
+            char_pos = 0
+            original_text = text
+
+            for token in tokens:
+                if token in self.all_special_tokens:
+                    offset_mapping.append((0, 0))
+                else:
+                    token_text = token
+                    if token_text.startswith("##"):
+                        token_text = token_text[2:]
+                    elif token_text in ["Ġ", "Ċ", "ĉ"]:
+                        token_text = {"Ġ": " ", "Ċ": "\n", "ĉ": "\t"}.get(token_text, token_text)
+
+                    start = original_text.find(token_text, char_pos)
+                    if start == -1:
+                        start = char_pos
+                    end = start + len(token_text)
+                    offset_mapping.append((start, end))
+                    char_pos = end
+
+        # Truncation
+        if truncation and max_length and len(input_ids) > max_length:
+            input_ids = input_ids[:max_length]
+            attention_mask = attention_mask[:max_length]
+            if offset_mapping:
+                offset_mapping = offset_mapping[:max_length]
+
+        # Padding
+        if padding and max_length and len(input_ids) < max_length:
+            pad_length = max_length - len(input_ids)
+            input_ids = input_ids + [self.pad_token_id] * pad_length
+            attention_mask = attention_mask + [0] * pad_length
+            if offset_mapping:
+                offset_mapping = offset_mapping + [(0, 0)] * pad_length
+
+        result = {"input_ids": input_ids, "attention_mask": attention_mask}
+
+        if return_offsets_mapping:
+            result["offset_mapping"] = offset_mapping
+
+        if return_token_type_ids:
+            result["token_type_ids"] = [0] * len(input_ids)
+
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            import torch
+
+            result = {
+                k: (
+                    torch.tensor([v], dtype=torch.long)
+                    if k != "offset_mapping"
+                    else [torch.tensor(v, dtype=torch.long)]
+                )
+                for k, v in result.items()
+            }
+        elif return_tensors == "np":
+            import numpy as np
+
+            result = {
+                k: np.array([v]) if k != "offset_mapping" else [np.array(v)]
+                for k, v in result.items()
+            }
+
+        return result
+
+    def decode(
+        self,
+        token_ids: Union[int, List[int]],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = True,
+    ) -> str:
+        """Decode token IDs back to text."""
+        if hasattr(token_ids, "cpu"):
+            token_ids = token_ids.cpu().tolist()
+        elif isinstance(token_ids, int):
+            token_ids = [token_ids]
+
+        tokens = self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens)
+        text = self.convert_tokens_to_string(tokens)
+
+        if clean_up_tokenization_spaces:
+            text = " ".join(text.split())
+
+        return text
+
+    def batch_encode_plus(
+        self,
+        batch_text_or_pairs: List[str],
+        add_special_tokens: bool = False,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        padding: bool = False,
+        return_tensors: Optional[str] = None,
+        return_attention_mask: bool = True,
+        return_offsets_mapping: bool = False,
+        return_token_type_ids: bool = False,
+    ) -> Dict:
+        """Encode a batch of texts."""
+        if isinstance(batch_text_or_pairs, str):
+            batch_text_or_pairs = [batch_text_or_pairs]
+
+        batch_output = []
+        for text in batch_text_or_pairs:
+            encoding = self.encode(
+                text,
+                add_special_tokens=add_special_tokens,
+                max_length=max_length,
+                truncation=truncation,
+                padding=False,
+                return_tensors=None,
+                return_attention_mask=return_attention_mask,
+                return_offsets_mapping=return_offsets_mapping,
+                return_token_type_ids=return_token_type_ids,
+            )
+            batch_output.append(encoding)
+
+        # Apply padding across batch
+        if padding:
+            max_len = max(len(enc["input_ids"]) for enc in batch_output)
+            for enc in batch_output:
+                pad_length = max_len - len(enc["input_ids"])
+                if pad_length > 0:
+                    enc["input_ids"] = enc["input_ids"] + [self.pad_token_id] * pad_length
+                    enc["attention_mask"] = enc["attention_mask"] + [0] * pad_length
+                    if return_offsets_mapping and "offset_mapping" in enc:
+                        enc["offset_mapping"] = enc["offset_mapping"] + [(0, 0)] * pad_length
+                    if return_token_type_ids and "token_type_ids" in enc:
+                        enc["token_type_ids"] = enc["token_type_ids"] + [0] * pad_length
+
+        batch_result = {
+            "input_ids": [enc["input_ids"] for enc in batch_output],
+            "attention_mask": [enc["attention_mask"] for enc in batch_output],
+        }
+
+        if return_offsets_mapping:
+            batch_result["offset_mapping"] = [enc.get("offset_mapping", []) for enc in batch_output]
+        if return_token_type_ids:
+            batch_result["token_type_ids"] = [enc.get("token_type_ids", []) for enc in batch_output]
+
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            import torch
+
+            batch_result = {
+                "input_ids": torch.tensor(batch_result["input_ids"], dtype=torch.long),
+                "attention_mask": torch.tensor(batch_result["attention_mask"], dtype=torch.long),
+            }
+            if return_offsets_mapping:
+                batch_result["offset_mapping"] = [
+                    torch.tensor(m, dtype=torch.long) for m in batch_result["offset_mapping"]
+                ]
+            if return_token_type_ids:
+                batch_result["token_type_ids"] = torch.tensor(
+                    batch_result["token_type_ids"], dtype=torch.long
+                )
+        elif return_tensors == "np":
+            import numpy as np
+
+            batch_result = {
+                "input_ids": np.array(batch_result["input_ids"]),
+                "attention_mask": np.array(batch_result["attention_mask"]),
+            }
+            if return_offsets_mapping:
+                batch_result["offset_mapping"] = [
+                    np.array(m) for m in batch_result["offset_mapping"]
+                ]
+            if return_token_type_ids:
+                batch_result["token_type_ids"] = np.array(batch_result["token_type_ids"])
+
+        return batch_result
+
+    def __call__(
+        self,
+        text: str,
+        add_special_tokens: bool = False,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        padding: bool = False,
+        return_tensors: Optional[str] = None,
+        return_attention_mask: bool = True,
+        return_offsets_mapping: bool = True,
+        return_token_type_ids: bool = False,
+    ) -> Dict:
+        """Call method for easy encoding."""
+        return self.encode(
+            text,
+            add_special_tokens=add_special_tokens,
+            max_length=max_length,
+            truncation=truncation,
+            padding=padding,
+            return_tensors=return_tensors,
+            return_attention_mask=return_attention_mask,
+            return_offsets_mapping=return_offsets_mapping,
+            return_token_type_ids=return_token_type_ids,
+        )
+
+    def get_vocab(self) -> Dict[str, int]:
+        """Get the vocabulary."""
+        return self.vocab
+
+    def save_pretrained(self, save_directory: str):
+        """Save tokenizer to directory."""
+        os.makedirs(save_directory, exist_ok=True)
+
+        tokenizer_config = {
+            "vocab": self.vocab,
+            "special_tokens_map": self.special_tokens_map,
+            "tokenizer_type": self.tokenizer_type,
+            "merges": self.merges,
+            "pad_token": self.pad_token,
+            "unk_token": self.unk_token,
+            "cls_token": self.cls_token,
+            "sep_token": self.sep_token,
+            "mask_token": self.mask_token,
+            "bos_token": self.bos_token,
+            "eos_token": self.eos_token,
+            "pad_token_id": self.pad_token_id,
+            "unk_token_id": self.unk_token_id,
+            "cls_token_id": self.cls_token_id,
+            "sep_token_id": self.sep_token_id,
+            "mask_token_id": self.mask_token_id,
+            "bos_token_id": self.bos_token_id,
+            "eos_token_id": self.eos_token_id,
+        }
+
+        # Save config files
+        with open(
+            os.path.join(save_directory, "tokenizer_config.json"), "w", encoding="utf-8"
+        ) as f:
+            json.dump(tokenizer_config, f, ensure_ascii=False, indent=2)
+
+        with open(os.path.join(save_directory, "vocab.json"), "w", encoding="utf-8") as f:
+            json.dump(self.vocab, f, ensure_ascii=False, indent=2)
+
+        with open(
+            os.path.join(save_directory, "special_tokens_map.json"), "w", encoding="utf-8"
+        ) as f:
+            json.dump(self.special_tokens_map, f, ensure_ascii=False, indent=2)
+
+        if self.merges:
+            with open(os.path.join(save_directory, "merges.txt"), "w", encoding="utf-8") as f:
+                for merge in self.merges:
+                    f.write(merge + "\n")
+
+    @classmethod
+    def from_pretrained(
+        cls, save_directory: str, tokenizer_type: str = "wordpiece", adding_new_token_path: str = "",add_tab_newline_vocab: bool=False
+    ):
+        """Load tokenizer from directory."""
+        with open(os.path.join(save_directory, "vocab.json"), "r", encoding="utf-8") as f:
+            vocab = json.load(f)
+
+        special_tokens_map = {}
+        special_tokens_path = os.path.join(save_directory, "special_tokens_map.json")
+        if os.path.exists(special_tokens_path):
+            with open(special_tokens_path, "r", encoding="utf-8") as f:
+                special_tokens_map = json.load(f)
+
+        # Add new tokens if provided
+        if adding_new_token_path and os.path.exists(adding_new_token_path):
+            with open(adding_new_token_path, "r", encoding="utf-8") as f:
+                adding_new_tokens = f.read().split("\n")
+
+            updated = False
+            c = 0
+            for ch in adding_new_tokens:
+                if ch not in vocab and ch != " ":
+                    vocab[ch] = len(vocab)
+                    updated = True
+                    c += 1
+
+            if updated:
+                with open(os.path.join(save_directory, "vocab.json"), "w", encoding="utf-8") as f:
+                    json.dump(vocab, f, ensure_ascii=False, indent=2)
+                print(f"Updating vocab with {c} new tokens")
+
+        # Load merges for BPE
+        merges = []
+        merges_path = os.path.join(save_directory, "merges.txt")
+        if os.path.exists(merges_path):
+            with open(merges_path, "r", encoding="utf-8") as f:
+                merges = [line.strip() for line in f.readlines()]
+
+        tokenizer = cls(
+            vocab = vocab,
+            special_tokens_map = special_tokens_map,
+            tokenizer_type = tokenizer_type,
+            add_tab_newline_vocab = add_tab_newline_vocab,
+            
+        )
+        tokenizer.merges = merges
+
+        return tokenizer
+
+    def add_tokens(self, new_tokens: Union[str, List[str]]) -> int:
+        """Add new tokens to the vocabulary."""
+        if isinstance(new_tokens, str):
+            new_tokens = [new_tokens]
+
+        added_count = 0
+        for token in new_tokens:
+            if token not in self.vocab:
+                self.vocab[token] = len(self.vocab)
+                self.inv_vocab[len(self.vocab) - 1] = token
+                added_count += 1
+        return added_count
+
+    def add_special_tokens(self, special_tokens_dict: Dict[str, str]) -> int:
+        """Add special tokens to the tokenizer."""
+        added_count = 0
+        for key, token in special_tokens_dict.items():
+            if hasattr(self, key):
+                setattr(self, key, token)
+                if token not in self.vocab:
+                    self.vocab[token] = len(self.vocab)
+                    self.inv_vocab[len(self.vocab) - 1] = token
+                    added_count += 1
+
+        self.all_special_tokens = [
+            self.pad_token,
+            self.unk_token,
+            self.cls_token,
+            self.sep_token,
+            self.mask_token,
+            self.bos_token,
+            self.eos_token,
+        ]
+        self.all_special_ids = [
+            self.pad_token_id,
+            self.unk_token_id,
+            self.cls_token_id,
+            self.sep_token_id,
+            self.mask_token_id,
+            self.bos_token_id,
+            self.eos_token_id,
+        ]
+
+        return added_count
+
+    def __len__(self) -> int:
+        """Return vocabulary size."""
+        return len(self.vocab)

src/dytr/training/__init__.py

@@ -0,0 +1,41 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Training modules for Dynamic Transformers.
+
+This module provides the trainer class, dataset handling utilities,
+and data processing functions for multi-task training with continual learning support.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+
+from dytr.training.data import MultiTaskDataset, TaskAwareBatchSampler, collate_fn
+from dytr.training.dataset import SingleDatasetProcessing
+from dytr.training.trainer import Trainer
+from dytr.training.utils import process_qa_dataset, set_seed
+
+__all__ = [
+    "Trainer",
+    "SingleDatasetProcessing",
+    "MultiTaskDataset",
+    "TaskAwareBatchSampler",
+    "collate_fn",
+    "set_seed",
+    "process_qa_dataset",
+]
+

src/dytr/training/data.py

@@ -0,0 +1,217 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Data utilities for multi-task training.
+
+This module provides dataset classes and samplers for handling multiple tasks
+during training.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+
+import bisect
+from collections import defaultdict
+from typing import  Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+
+def collate_fn(batch, fixed_max_len: Optional[int] = None):
+    """
+    Custom collate function for multi-task batches.
+
+    Handles variable-length sequences and different task types.
+
+    Args:
+        batch: List of samples from the dataset
+
+    Returns:
+        Dictionary with batched tensors
+    """
+    if not isinstance(batch, list):
+        batch = [batch]
+
+    batch_dict = {}
+
+    for key in batch[0].keys():
+        if key in ["input_ids", "attention_mask", "labels"]:
+            values = []
+            for item in batch:
+                if key in item and isinstance(item[key], torch.Tensor):
+                    if item[key].dim() > 0:
+                        values.append(item[key])
+
+            if not values:
+                continue
+
+            # Find max length for padding
+            if fixed_max_len is not None:
+                # Use fixed length for all sequences
+                max_len = batch[0]["max_length"]  # fixed_max_len 
+            else:
+                # Dynamic: use max length in this batch
+                max_len = max(v.size(0) for v in values if v.dim() > 0)
+            # max_len = max(v.size(0) for v in values if v.dim() > 0)
+            padded_values = []
+
+            for v in values:
+                if v.dim() > 0:
+                    if v.size(0) < max_len:
+                        pad_size = max_len - v.size(0)
+                        if key == "labels":
+                            padded = F.pad(v, (0, pad_size), value=-100)
+                        else:
+                            padded = F.pad(v, (0, pad_size), value=0)
+                    else:
+                        padded = v
+                    padded_values.append(padded.unsqueeze(0))
+
+            if padded_values:
+                batch_dict[key] = torch.cat(padded_values, dim=0)
+        else:
+            # For non-tensor fields, keep as list
+            batch_dict[key] = [item[key] for item in batch if key in item]
+
+    return batch_dict
+
+
+class TaskAwareBatchSampler:
+    """
+    Batch sampler that creates batches with samples from the same task.
+
+    Ensures that each batch contains samples from only one task to simplify
+    task-specific processing.
+
+    Args:
+        dataset: MultiTaskDataset instance
+        batch_size: Batch size
+        shuffle: Whether to shuffle batches
+        drop_last: Whether to drop incomplete batches
+    """
+
+    def __init__(self, dataset, batch_size, shuffle=True, drop_last=True):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+
+        # Group indices by task
+        self.indices_by_task = defaultdict(list)
+        for idx in range(len(dataset)):
+            dataset_idx = bisect.bisect_right(dataset.cumulative_sizes, idx) - 1
+            task_name = dataset.task_names[dataset_idx]
+            self.indices_by_task[task_name].append(idx)
+
+        # Create batches
+        self.batches = []
+        self.batch_task_map = []
+
+        for task_name, indices in self.indices_by_task.items():
+            if shuffle:
+                np.random.shuffle(indices)
+
+            for i in range(0, len(indices), batch_size):
+                batch_indices = indices[i : i + batch_size]
+                if drop_last and len(batch_indices) < batch_size:
+                    continue
+                self.batches.append(batch_indices)
+                self.batch_task_map.append(task_name)
+
+        # Shuffle batches if needed
+        if shuffle:
+            combined = list(zip(self.batches, self.batch_task_map))
+            np.random.shuffle(combined)
+            if combined:
+                self.batches, self.batch_task_map = zip(*combined)
+                self.batches = list(self.batches)
+                self.batch_task_map = list(self.batch_task_map)
+
+    def __iter__(self):
+        """Iterate over batches."""
+        for batch_indices in self.batches:
+            yield batch_indices
+
+    def __len__(self):
+        """Number of batches."""
+        return len(self.batches)
+
+
+class MultiTaskDataset(Dataset):
+    """
+    Dataset that combines multiple single-task datasets.
+
+    Creates a unified dataset from multiple task-specific datasets,
+    allowing seamless multi-task training.
+
+    Args:
+        datasets_dict: Dictionary mapping task names to (dataset, strategy) tuples
+        tokenizer: Tokenizer instance
+        task_configs: List of task configurations
+    """
+
+    def __init__(self, datasets_dict, tokenizer, task_configs):
+        self.datasets = []
+        self.task_names = []
+        self.max_lengths ={t:datasets_dict[t][0].max_len for t in datasets_dict.keys() }
+        
+        #self.max_lengths = {t.task_name: t.max_length for t in task_configs}
+        self.strategies = []
+        self.cumulative_sizes = [0]
+
+        print("Max lengths per task:", self.max_lengths)
+
+        for task_name, (dataset, strategy) in datasets_dict.items():
+            if len(dataset) > 0:
+                self.datasets.append(dataset)
+                self.task_names.append(task_name)
+                self.strategies.append(strategy)
+                self.cumulative_sizes.append(self.cumulative_sizes[-1] + len(dataset))
+
+    def __len__(self):
+        """Total number of samples across all tasks."""
+        return self.cumulative_sizes[-1]
+
+    def __getitem__(self, idx):
+        """
+        Get a sample by index.
+
+        Args:
+            idx: Index in the combined dataset
+
+        Returns:
+            Sample with task_name and strategy added
+        """
+        if idx < 0:
+            idx = len(self) + idx
+
+        # Find which dataset this index belongs to
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) - 1
+        sample_idx = idx - self.cumulative_sizes[dataset_idx]
+
+        # Get the sample from the specific dataset
+        item = self.datasets[dataset_idx][sample_idx]
+
+        if isinstance(item, dict):
+            item["task_name"] = self.task_names[dataset_idx]
+            item["strategy"] = self.strategies[dataset_idx].value
+            item["max_length"] = self.max_lengths[self.task_names[dataset_idx]]
+
+        return item

src/dytr/training/dataset.py

@@ -0,0 +1,413 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Dataset classes for single-task data processing.
+
+This module provides the SingleDatasetProcessing class which handles
+dataset preparation for all training strategies.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+
+import hashlib
+import pickle
+from pathlib import Path
+#from typing import Optional
+
+#import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+from dytr.core.config import TrainingStrategy
+
+def process_tags_column(df, tags_column='tags',label_to_ids={},calculate_distribution=False):
+    """
+    Process tags column: convert string tags to numeric labels.
+    
+    Args:
+        df: DataFrame with tags column
+        tags_column: Name of tags column (default: 'tags')
+        
+    Returns:
+        df: DataFrame with processed tags column
+        label_to_ids: Dictionary mapping tag labels to IDs
+        distribution: Dictionary of tag distribution
+    """
+    from collections import Counter
+    df[tags_column] = df[tags_column].apply(lambda x: ' '.join(map(str, x)) if isinstance(x, list) else x)
+    # Collect all unique tags
+    all_tags = set()
+    for tags_str in df[tags_column].dropna():
+        if isinstance(tags_str, str):
+            tags = tags_str.split()
+            all_tags.update(tags)
+    
+    if not label_to_ids:
+        # Create label mapping (O -> 0, others start from 1)
+        
+        label_to_ids = {'O': 0}
+        
+        for tag in sorted(all_tags):
+            if tag != 'O' and not tag.isdigit():
+                label_to_ids[tag] = len(label_to_ids)
+            elif tag.isdigit() and int(tag) !=0:
+                label_to_ids[tag] = int(tag)#len(label_to_ids)
+    
+    
+    # Convert tags to numeric labels
+    def convert_tags(tags_str):
+        if not isinstance(tags_str, str):
+            return ''
+        tags = tags_str.split()
+        return ' '.join(str(label_to_ids.get(tag, 0)) for tag in tags)
+    
+    df[tags_column] = df[tags_column].apply(convert_tags)
+    if not calculate_distribution:
+        return df, label_to_ids, {}
+    # Calculate distribution
+    distribution = Counter()
+    for tags_str in df[tags_column]:
+        if tags_str:
+            for label in map(int, tags_str.split()):
+                distribution[label] += 1
+    
+    # Convert to readable format
+    readable_dist = {label: count for label, count in distribution.items()}
+    
+    return df, label_to_ids, readable_dist
+
+class SingleDatasetProcessing(Dataset):
+    """
+    Dataset class for processing single tasks with different training strategies.
+
+    Supports:
+    - Sentence Classification
+    - Token Classification
+    - Seq2Seq
+    - Causal Language Modeling
+
+    Args:
+        df: DataFrame containing the data
+        tokenizer: Tokenizer instance
+        max_len: Maximum sequence length
+        task_name: Name of the task
+        strategy: Training strategy
+        num_labels: Number of labels (for classification)
+        text_column: Column name for text
+        label_column: Column name for labels
+        source_column: Column name for source text (seq2seq)
+        target_column: Column name for target text (seq2seq)
+        tags_column: Column name for token tags
+        stride: Stride for windowing (causal LM)
+        cache_dir: Directory for caching processed data
+    """
+
+    def __init__(
+        self,
+        df,
+        tokenizer,
+        max_len,
+        task_name,
+        strategy,
+        num_labels=None,
+        text_column="text",
+        label_column="label",
+        source_column="source",
+        target_column="target",
+        tags_column="tags",
+        stride=None,
+        cache_dir="./dataset_cache",
+        token_labeling_first_only=True,
+        label_to_ids={}
+    ):
+
+        self.df = df
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+        self.task_name = task_name
+        self.strategy = strategy
+        self.num_labels = num_labels
+        self.text_column = text_column
+        self.stride = stride if stride is not None else max_len // 2
+        self.samples = []
+        self.pad_id = tokenizer.pad_token_id or 0
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.token_labeling_first_only = token_labeling_first_only
+        self.label_to_ids=label_to_ids
+        # Process based on strategy
+        if strategy == TrainingStrategy.SENTENCE_CLASSIFICATION:
+            self._process_classification(label_column)
+        elif strategy == TrainingStrategy.TOKEN_CLASSIFICATION:
+            self._process_token_classification(tags_column)
+        elif strategy == TrainingStrategy.SEQ2SEQ:
+            self._process_seq2seq(source_column, target_column)
+        elif strategy == TrainingStrategy.CAUSAL_LM:
+            self._process_causal_lm(text_column)
+
+        print(f"  Created {len(self.samples)} samples for {self.task_name} >> {strategy.value}")
+
+    def _process_classification(self, label_column):
+        """Process data for sentence classification."""
+        if self.num_labels is None:
+            self.num_labels = len(self.df[label_column].unique())
+
+        for _, row in self.df.iterrows():
+            label = int(row[label_column])
+            self.samples.append({"text": row[self.text_column], "label": label})
+            if len(self.samples) == 1:
+                print(f"    Sample text: {str(row[self.text_column])}...")
+                print(f"    Sample label: {label}")
+
+    def _process_token_classification(self, tags_column):
+        """Process data for token classification."""
+        if not self.label_to_ids:
+            self.df, self.label_to_ids, _=process_tags_column(self.df, tags_column)
+        else:
+            self.df, _, _=process_tags_column(self.df, tags_column,label_to_ids=self.label_to_ids)
+        """
+        all_tags = set()
+        for tags in self.df[tags_column]:
+            if isinstance(tags, str):
+                for tag in tags.split():
+                    all_tags.add(int(tag))
+            elif isinstance(tags, list):
+                for tag in tags:
+                    all_tags.add(int(tag))
+        all_tags = sorted(all_tags)"""
+        self.num_labels = max(self.label_to_ids.values())+1 #len(all_tags)
+
+        valid_samples = 0
+        for _, row in self.df.iterrows():
+            text = row[self.text_column]
+            if isinstance(row[tags_column], str):
+                tags = [int(i) for i in row[tags_column].split()]
+            else:
+                tags = [int(i) for i in row[tags_column]]
+
+            words = text.split()
+            if len(tags) == len(words):
+                self.samples.append({"text": text, "tags": tags})
+                valid_samples += 1
+                if valid_samples == 1:
+                    print(f"    Sample text: {text}...")
+                    print(f"    Sample tags: {tags}...")
+
+    def _process_seq2seq(self, source_column, target_column):
+        """Process data for seq2seq."""
+        self.num_labels=len(self.tokenizer)
+        for _, row in self.df.iterrows():
+            self.samples.append({"source": row[source_column], "target": row[target_column]})
+            if len(self.samples) == 1:
+                print(f"    Sample source: {str(row[source_column])}...")
+                print(f"    Sample target: {str(row[target_column])}...")
+
+    def _process_causal_lm(self, text_column):
+        """Process data for causal language modeling with windowing."""
+        self.num_labels=len(self.tokenizer)
+        cache_key = hashlib.md5(
+            f"{len(self.df)}_{self.max_len}_{self.stride}_{text_column}_{self.tokenizer.name_or_path}".encode()
+        ).hexdigest()
+        cache_file = self.cache_dir / f"causal_windows_{cache_key}.pkl"
+
+        if cache_file.exists():
+            print(f"  Loading cached window texts from {cache_file}")
+            with open(cache_file, "rb") as f:
+                cache_data = pickle.load(f)
+                self.samples = cache_data["window_texts"]
+                print(f"    Loaded {len(self.samples)} cached windows")
+        else:
+            print("  Creating causal LM dataset with dynamic window sampling...")
+
+            if isinstance(self.df, pd.DataFrame):
+                texts = self.df[text_column].astype(str).tolist()
+            else:
+                texts = [t[text_column] for t in self.df]
+
+            for doc_idx, text in enumerate(tqdm(texts, desc="Processing documents")):
+                tokens = self.tokenizer.tokenize(text)
+                token_count = len(tokens)
+
+                if token_count <= self.max_len:
+                    self.samples.append(text)
+                else:
+                    num_windows = (token_count - self.max_len) // self.stride + 1
+                    for window_idx in range(num_windows):
+                        start_token = window_idx * self.stride
+                        end_token = min(start_token + self.max_len, token_count)
+
+                        window_tokens = tokens[start_token:end_token]
+                        if len(window_tokens) < 2:
+                            continue
+                        window_text = self.tokenizer.convert_tokens_to_string(window_tokens)
+                        self.samples.append(window_text)
+
+            # Cache the processed windows
+            cache_data = {"window_texts": self.samples}
+            response=input("Do you want to save Causal Data Window in cach dir?[Y/n]")
+            if response=="Y":
+                with open(cache_file, "wb") as f:
+                    pickle.dump(cache_data, f)
+                print(f"data saved on dir: {cache_file}")
+
+            print(f"    Created and cached {len(self.samples)} windows from {len(texts)} documents")
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        """Get a single sample based on the strategy."""
+        if self.strategy == TrainingStrategy.SENTENCE_CLASSIFICATION:
+            return self._get_classification_item(idx)
+        elif self.strategy == TrainingStrategy.TOKEN_CLASSIFICATION:
+            return self._get_token_classification_item(idx)
+        elif self.strategy == TrainingStrategy.SEQ2SEQ:
+            return self._get_seq2seq_item(idx)
+        elif self.strategy == TrainingStrategy.CAUSAL_LM:
+            return self._get_causal_lm_item(idx)
+
+    def _get_classification_item(self, idx):
+        """Get item for sentence classification."""
+        sample = self.samples[idx]
+        encoding = self.tokenizer(
+            sample["text"],
+            max_length=self.max_len,
+            truncation=True,
+            padding=False,
+            add_special_tokens=False,
+        )
+        input_ids = encoding["input_ids"]
+        if not input_ids:
+            input_ids = [self.tokenizer.pad_token_id or 0]
+        attention_mask = [1] * len(input_ids)
+        labels = torch.full((len(input_ids),), -100, dtype=torch.long)
+        labels[0] = sample["label"]
+
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "labels": labels,
+            "task_name": self.task_name,
+            "strategy": self.strategy.value,
+        }
+
+    def _get_token_classification_item(self, idx):
+        """Get item for token classification."""
+        sample = self.samples[idx]
+        encoding = self.tokenizer(
+            sample["text"],
+            max_length=self.max_len,
+            truncation=True,
+            padding=False,
+            return_offsets_mapping=True,
+            add_special_tokens=False,
+        )
+        input_ids = encoding["input_ids"]
+        if not input_ids:
+            input_ids = [self.tokenizer.pad_token_id or 0]
+        attention_mask = [1] * len(input_ids)
+
+        # Get word ids for token alignment
+        try:
+            word_ids = self.tokenizer.word_ids(encoding)
+        except:
+            word_ids = encoding.word_ids()
+
+        labels = [-100] * len(input_ids)
+        if self.token_labeling_first_only:
+            previous_word_idx = None
+            for i, word_idx in enumerate(word_ids):
+                if word_idx is None:
+                    continue
+                if word_idx != previous_word_idx:
+                    if word_idx < len(sample["tags"]):
+                        tag = sample["tags"][word_idx]
+                        label_id = int(tag)
+                        if label_id < self.num_labels:
+                            labels[i] = label_id
+                    previous_word_idx = word_idx
+        else:
+            
+            # Label every token
+            for i, word_idx in enumerate(word_ids):
+                if word_idx is None:
+                    continue
+                if word_idx < len(sample["tags"]):
+                    tag = sample["tags"][word_idx]
+                    label_id = int(tag)
+                    if label_id < self.num_labels:
+                        labels[i] = label_id
+
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+            "task_name": self.task_name,
+            "strategy": self.strategy.value,
+        }
+
+    def _get_seq2seq_item(self, idx):
+        """Get item for seq2seq."""
+        sample = self.samples[idx]
+        source_encoding = self.tokenizer(
+            sample["source"],
+            max_length=self.max_len,
+            truncation=True,
+            padding=False,
+            add_special_tokens=False,
+        )
+        target_encoding = self.tokenizer(
+            sample["target"],
+            max_length=self.max_len,
+            truncation=True,
+            padding=False,
+            add_special_tokens=False,
+        )
+
+        return {
+            "input_ids": torch.tensor(source_encoding["input_ids"], dtype=torch.long),
+            "labels": torch.tensor(target_encoding["input_ids"], dtype=torch.long),
+            "attention_mask": torch.ones(len(source_encoding["input_ids"]), dtype=torch.long),
+            "task_name": self.task_name,
+            "strategy": self.strategy.value,
+        }
+
+    def _get_causal_lm_item(self, idx):
+        """Get item for causal language modeling."""
+        window_text = self.samples[idx]
+        encoding = self.tokenizer(
+            window_text,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=self.max_len,
+            padding=False,
+        )
+        input_ids = encoding["input_ids"]
+
+        # Shift labels for next token prediction
+        labels = input_ids[1:] + [self.pad_id]
+
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+            "attention_mask": torch.ones(len(input_ids), dtype=torch.long),
+            "task_name": self.task_name,
+            "strategy": self.strategy.value,
+        }

src/dytr/training/trainer.py

@@ -0,0 +1,572 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Trainer class for multi-task continual learning.
+
+This module provides the Trainer class which handles training,
+validation, and continual learning features like EWC and experience replay.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+
+import math
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+#import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from dytr.core.config import ModelConfig, TaskConfig
+from dytr.core.model import DynamicTransformer
+from dytr.training.data import MultiTaskDataset, TaskAwareBatchSampler, collate_fn
+from dytr import EWC
+from dytr.utils.logging import get_logger
+
+class Trainer:
+    """
+    Trainer class for multi-task continual learning.
+
+    Handles training across multiple tasks with support for:
+    - EWC (Elastic Weight Consolidation)
+    - Experience replay
+    - Learning rate scheduling
+    - Gradient accumulation
+    - Early stopping
+    - Model checkpointing
+
+    Args:
+        model: DynamicTransformer model
+        config: Model configuration
+        exp_dir: Directory for saving experiment outputs
+    """
+
+    def __init__(self, model: DynamicTransformer, config: ModelConfig, exp_dir: str):
+        self.logger = get_logger(__name__)
+        self.model = model
+        self.config = config
+        self.device = config.device
+        self.exp_dir = Path(exp_dir)
+        self.exp_dir.mkdir(parents=True, exist_ok=True)
+        self.best_val_loss = model.best_val_loss
+        self.best_metrics = {}
+        self.patience_counter = 0
+        self.loss_history = []
+        self.num_labels_per_task={}
+        self.val_loss_history = []
+        self.metrics_history = []
+        self.logger.info("Trainer initialized")
+        #self.logger.debug(f"Experiment directory: {exp_dir}")
+        self.logger.debug("Device: %s ",self.config.device)
+
+    def count_parameters(self):
+        """Count and display model parameters."""
+        total = sum(p.numel() for p in self.model.parameters())
+        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+
+        self.logger.info("  Model size: %s parameters",total)
+        self.logger.info("  Trainable: %s" ,trainable)
+        self.logger.info("\nParameter breakdown:")
+
+        for name, module in self.model.task_heads.items():
+            self.logger.info("    Head %s: %s",name, sum(p.numel() for p in module.parameters()))
+
+        for name, module in self.model.decoders.items():
+            self.logger.info("    Decoder %s: %s",name, sum(p.numel() for p in module.parameters()))
+
+        self.logger.info(f"    Shared encoder: {sum(p.numel() for p in self.model.encoder.parameters()):,}")
+        self.logger.info(
+            f"    Shared embedding: {sum(p.numel() for p in self.model.shared_embedding.parameters()):,}"
+        )
+
+        return total
+    def freeze_all(self):
+        """
+        Freeze all parameters in the model.
+        Used as initial state before selectively unfreezing tasks.
+        """
+        for param in self.model.parameters():
+            param.requires_grad = False
+        self.logger.info("  Frozen all model parameters")
+    
+    def unfreeze_shared_encoder(self):
+        """
+        Unfreeze the shared encoder and embedding layers.
+        These are shared across all tasks and should be trainable for all tasks.
+        """
+        # Unfreeze shared embedding
+        for param in self.model.shared_embedding.parameters():
+            param.requires_grad = True
+        
+        # Unfreeze encoder
+        for param in self.model.encoder.parameters():
+            param.requires_grad = True
+        
+        self.logger.info("  Unfrozen shared encoder and embeddings")
+    
+    def unfreeze_task_components(self, task_name: str):
+        """
+        Unfreeze task-specific components for a given task.
+        
+        Args:
+            task_name: Name of the task to unfreeze
+        """
+        # Unfreeze task head if it exists
+        if task_name in self.model.task_heads:
+            for param in self.model.task_heads[task_name].parameters():
+                param.requires_grad = True
+            self.logger.info(f"  Unfrozen head for task: {task_name}")
+        
+        # Unfreeze decoder if it exists
+        if task_name in self.model.decoders:
+            for param in self.model.decoders[task_name].parameters():
+                param.requires_grad = True
+            self.logger.info(f"  Unfrozen decoder for task: {task_name}")
+    
+    def unfreeze_for_tasks(self, task_names: list, unfreeze_shared: bool = True):
+        """
+        Unfreeze model components for specific tasks.
+        
+        Args:
+            task_names: List of tasks to unfreeze
+            unfreeze_shared: Whether to unfreeze shared encoder (default: True)
+        """
+        # Unfreeze shared components first (if needed)
+        if unfreeze_shared:
+            self.unfreeze_shared_encoder()
+        
+        # Unfreeze each task's specific components
+        for task_name in task_names:
+            self.unfreeze_task_components(task_name)
+    
+    
+    def train(self, task_configs: List[TaskConfig], train_datasets: Dict, val_datasets: Dict):
+        """
+        Train the model on multiple tasks.
+
+        Args:
+            task_configs: List of task configurations
+            train_datasets: Dictionary of training datasets per task
+            val_datasets: Dictionary of validation datasets per task
+
+        Returns:
+            Trained model
+        """
+        self.logger.info("=" * 60)
+        self.logger.info("Starting training session")
+        self.logger.info("=" * 60)
+        training_task_names = list(train_datasets.keys())
+        encoder_unfreeze = not all(train_datasets[i][1] == "causal_lm" for i in train_datasets.keys())
+        
+        self.logger.debug(f"Task configs: {[tc.task_name for tc in task_configs]}")
+        self.logger.debug(f"Train datasets: {training_task_names}")
+        self.logger.debug(f"Val datasets: {list(val_datasets.keys())}")
+        # Add tasks to model
+        self.freeze_all()
+        for task_config in task_configs:
+            self.model.add_task(task_config)
+            if task_config.task_name not in self.num_labels_per_task:
+                self.num_labels_per_task[task_config.task_name]=task_config.num_labels
+        
+        self.unfreeze_for_tasks(training_task_names,encoder_unfreeze)
+        if len(train_datasets.keys())<1:
+            self.logger.error("No training datasets found for tasks!")
+            return self.model
+        
+        self.logger.info("Model parameters after adding tasks:")
+        self.count_parameters()
+
+        # Create multi-task datasets
+        train_dataset = MultiTaskDataset(train_datasets, self.model.tokenizer, task_configs)
+        val_dataset = MultiTaskDataset(val_datasets, self.model.tokenizer, task_configs)
+
+        # Create samplers and dataloaders
+        train_sampler = TaskAwareBatchSampler(
+            train_dataset, self.config.per_device_train_batch_size, shuffle=True
+        )
+        use_fixed_padding = self.config.use_replay and len(self.model.replay_buffer.buffer) > 0
+
+        def train_collate(batch):
+            if use_fixed_padding:
+                # Use fixed max length from config
+                return collate_fn(batch, fixed_max_len=self.config.max_seq_len)
+            else:
+                # Use dynamic padding
+                return collate_fn(batch)
+
+        train_loader = DataLoader(
+            train_dataset, batch_sampler=train_sampler, collate_fn=train_collate, num_workers=0
+        )
+        
+        if len(val_dataset) > 0:
+            val_sampler = TaskAwareBatchSampler(
+                val_dataset, self.config.per_device_eval_batch_size, shuffle=False
+            )
+            val_loader = DataLoader(
+                val_dataset, batch_sampler=val_sampler, collate_fn=collate_fn, num_workers=0
+            )
+        else:
+            val_loader = None
+            self.logger.warning("  Warning: No validation data provided")
+
+        # Setup optimizer with task-specific learning rates
+        param_groups = []
+        assigned = set()
+
+        # Task heads parameters
+        head_params = []
+        for th in self.model.task_heads.values():
+            for p in th.parameters():
+                if id(p) not in assigned:
+                    head_params.append(p)
+                    assigned.add(id(p))
+        if head_params:
+            param_groups.append(
+                {
+                    "params": head_params,
+                    "lr": self.config.learning_rate * self.config.head_lr_mult,
+                    "name": "task_heads",
+                }
+            )
+
+        # Decoder parameters
+        decoder_params = []
+        for d in self.model.decoders.values():
+            for p in d.parameters():
+                if id(p) not in assigned:
+                    decoder_params.append(p)
+                    assigned.add(id(p))
+        if decoder_params:
+            param_groups.append(
+                {
+                    "params": decoder_params,
+                    "lr": self.config.learning_rate * self.config.decoder_lr_mult,
+                    "name": "decoders",
+                }
+            )
+
+        # Shared parameters
+        shared_params = []
+        for p in self.model.parameters():
+            if id(p) not in assigned:
+                shared_params.append(p)
+                assigned.add(id(p))
+        if shared_params:
+            param_groups.append(
+                {
+                    "params": shared_params,
+                    "lr": self.config.learning_rate * self.config.shared_lr_mult,
+                    "name": "shared",
+                }
+            )
+
+        # Optimizer and scheduler
+        optimizer = optim.AdamW(param_groups, weight_decay=self.config.weight_decay)
+        total_steps = len(train_loader) * self.config.num_train_epochs
+
+        def lr_lambda(step):
+            if step < self.config.warmup_steps:
+                return step / max(1, self.config.warmup_steps)
+            progress = (step - self.config.warmup_steps) / max(
+                1, total_steps - self.config.warmup_steps
+            )
+            return max(
+                self.config.min_learning_rate / self.config.learning_rate,
+                0.5 * (1 + math.cos(math.pi * progress)),
+            )
+
+        scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+        best_model_path = os.path.join(self.exp_dir, "best_model.pt")
+        steps=0
+
+
+
+
+        # Training loop
+        for epoch in range(self.config.num_train_epochs):
+            self.model.train()
+            total_loss = 0
+            num_batches = 0
+            task_losses = defaultdict(list)
+            epoch_metrics = defaultdict(lambda: defaultdict(list))
+
+            pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{self.config.num_train_epochs}")
+
+            for batch_idx, batch in enumerate(pbar):
+                try:
+                    task_name = batch["task_name"][0]
+                    strategy=batch["strategy"][0]
+
+                    # Move batch to device
+                    
+                    batch = {
+                        k: v.to(self.device) if isinstance(v, torch.Tensor) else v
+                        for k, v in batch.items()
+                    }
+
+                    # Add replay samples if available
+                    if self.config.use_replay and len(self.model.replay_buffer.buffer) > 0:
+                        replay_samples = self.model.replay_buffer.sample(
+                            self.config.per_device_train_batch_size // 4
+                        )
+                        if replay_samples:
+                            
+                            for replay in replay_samples:
+                                if task_name in self.num_labels_per_task and replay["task_name"] in self.num_labels_per_task and self.num_labels_per_task[replay["task_name"]]!=self.num_labels_per_task[task_name]:
+                                    break
+                                elif task_name not in self.num_labels_per_task or replay["task_name"] not in self.num_labels_per_task:
+                                    break
+                                elif replay["strategy"]!=strategy:
+                                    break
+                                #print(replay["strategy"])
+                                for k, v in replay.items():
+                                    if isinstance(v, torch.Tensor):
+                                        if batch[k].dim() == v.dim():
+                                            # Same dimensions, can concatenate
+                                            batch[k] = torch.cat(
+                                                [batch[k], v.to(self.device)], dim=0
+                                            )
+                                        else:
+                                            # Different dimensions, need to handle specially
+                                            # print(f"  Warning: Cannot concatenate {k} with dim {batch[k].dim()} and {v.dim()}")
+                                            # Skip this replay sample for this key
+                                            pass
+                                        # batch[k] = torch.cat([batch[k], v.to(self.device)], dim=0)
+
+                    # Forward pass
+                    outputs = self.model(
+                        input_ids=batch["input_ids"],
+                        attention_mask=batch.get("attention_mask"),
+                        task_name=task_name,
+                        labels=batch.get("labels"),
+                    )
+
+                    if outputs is not None and "loss" in outputs and outputs["loss"] is not None:
+                        loss = outputs["loss"]
+
+                        # Add EWC penalty if enabled
+                        if self.config.use_ewc and len(self.model.ewc_penalties) > 0:
+                            ewc_loss = 0
+                            for ts in self.model.ewc_penalties.keys():
+                                if ts in training_task_names:
+                                    continue
+                                ewc_loss += self.model.ewc_penalties[ts].penalty(self.model)
+                                
+                            #for ewc in self.model.ewc_penalties.values():
+                            #    ewc_loss += ewc.penalty(self.model)
+                            loss = loss + self.config.ewc_lambda * ewc_loss
+
+                        if torch.isfinite(loss):
+                            loss.backward()
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.parameters(), self.config.max_grad_norm
+                            )
+                            optimizer.step()
+                            scheduler.step()
+                            optimizer.zero_grad()
+
+                            total_loss += loss.item()
+                            num_batches += 1
+                            steps+=1
+                            self.loss_history.append(loss.item())
+                            task_losses[task_name].append(loss.item())
+
+                            # Collect metrics
+                            if "metrics" in outputs and outputs["metrics"]:
+                                for metric_name, value in outputs["metrics"].items():
+                                    if isinstance(value, (int, float)):
+                                        epoch_metrics[task_name][metric_name].append(value)
+                                    elif isinstance(value, torch.Tensor):
+                                        epoch_metrics[task_name][metric_name].append(value.item())
+                        else:
+                            self.logger.warning(f"  Warning: Non-finite loss at batch {batch_idx}: {loss}")
+
+                    # Logging
+                    if steps % self.config.logging_steps == 0 and num_batches > 0:
+                        task_avg = {
+                            task: sum(losses[-100:]) / len(losses[-100:])
+                            for task, losses in task_losses.items()
+                            if losses
+                        }
+                        pbar.set_postfix(
+                            {"loss": f"{total_loss/num_batches:.4f}", "tasks": str(task_avg)}
+                        )
+
+                    # Validation during training (if steps strategy)
+                    if (
+                        self.config.evaluation_strategy == "steps"
+                        and steps % self.config.validation_check_interval == 0
+                        and num_batches > 0
+                        and val_loader is not None
+                    ):
+                        val_results = self.validate(val_loader)
+                        if val_results:
+                            self._handle_validation(val_results, best_model_path)
+
+                except RuntimeError as e:
+                    if "out of memory" in str(e):
+                        if torch.cuda.is_available():
+                            torch.cuda.empty_cache()
+                        optimizer.zero_grad()
+                        self.logger.warning(f"  Error in training batch: {e}")
+                        continue
+                    else:
+                        raise e
+                except Exception as e:
+                    self.logger.warning(f"  Error in training batch: {e}")
+                    continue
+
+            # Epoch end
+            if num_batches > 0:
+                avg_epoch_loss = total_loss / num_batches
+                print(f"\n  Epoch {epoch+1} average training loss: {avg_epoch_loss:.4f}")
+
+                for task, metrics in epoch_metrics.items():
+                    print(f"    {task} training metrics:")
+                    for metric_name, values in metrics.items():
+                        if values:
+                            avg_value = sum(values) / len(values)
+                            print(f"      {metric_name}: {avg_value:.4f}")
+
+            # Validation at epoch end
+            if val_loader is not None and self.config.evaluation_strategy != "steps":
+                val_results = self.validate(val_loader)
+                if val_results:
+                    self._handle_validation(val_results, best_model_path)
+
+                    if self.patience_counter >= self.config.patience:
+                        print(f"  Early stopping triggered after {epoch+1} Validation")
+                        response=input("Insert letter Y to stop the training: [Y/n]")
+                        if response=='Y':
+                            break
+                        self.patience_counter=0
+                        
+                            
+        
+        self.logger.debug(f"\n  Training completed. Best validation loss: {self.best_val_loss:.4f}")
+        if self.config.use_ewc:# and use_fixed_padding:
+            for task_name in train_datasets.keys():
+                if task_name in self.model.ewc_penalties: continue
+                ewc = EWC(self.model, task_name, lambda_param=self.config.ewc_lambda)
+                fisher_loader = DataLoader(train_datasets[task_name][0], batch_size= 8, shuffle=True, collate_fn=collate_fn)
+                ewc.compute_fisher(fisher_loader, self.config.device)
+                self.model.ewc_penalties[task_name] = ewc
+        return self.model
+
+    def _handle_validation(self, val_results, best_model_path):
+        """Handle validation results and model checkpointing."""
+        val_loss = val_results["avg_loss"]
+        val_metrics = val_results["metrics"]
+
+        self.val_loss_history.append(val_loss)
+        self.metrics_history.append(val_metrics)
+        print(f"\n  Validation loss: {val_loss:.4f}")
+
+        for task, metrics in val_metrics.items():
+            print(f"    {task} validation:")
+            for metric_name, value in metrics.items():
+                if isinstance(value, dict):
+                    print(f"      {metric_name}:")
+                    for k, v in value.items():
+                        print(f"        {k}: {v:.4f}")
+                else:
+                    print(f"      {metric_name}: {value:.4f}")
+
+        if val_loss < self.best_val_loss:
+            improvement = self.best_val_loss - val_loss
+            self.model.steps_without_improvement = 0
+            self.best_val_loss = val_loss
+            self.model.best_val_loss = val_loss
+            self.best_metrics = val_metrics
+            self.patience_counter = 0
+            self.model.save_model(best_model_path)
+            self.logger.warning("  ✓ Best model saved improvement: %s",improvement)
+        else:
+            self.model.steps_without_improvement += 1
+            self.patience_counter += 1
+            self.logger.warning("  No improvement for %s epochs",self.patience_counter)
+
+    def validate(self, val_loader):
+        """
+        Validate the model on validation data.
+
+        Args:
+            val_loader: DataLoader for validation data
+
+        Returns:
+            Dictionary with average loss and metrics per task
+        """
+        self.model.eval()
+        total_loss = 0
+        val_batches = 0
+        task_metrics = defaultdict(lambda: defaultdict(list))
+
+        with torch.no_grad():
+            for batch in val_loader:
+                task_name = batch["task_name"][0]
+
+                batch = {
+                    k: v.to(self.device) if isinstance(v, torch.Tensor) else v
+                    for k, v in batch.items()
+                }
+
+                outputs = self.model(
+                    input_ids=batch["input_ids"],
+                    attention_mask=batch.get("attention_mask"),
+                    task_name=task_name,
+                    labels=batch.get("labels"),
+                )
+
+                if outputs is not None and "loss" in outputs and outputs["loss"] is not None:
+                    if torch.isfinite(outputs["loss"]):
+                        total_loss += outputs["loss"].item()
+                        val_batches += 1
+
+                        if "metrics" in outputs and outputs["metrics"]:
+                            for metric_name, value in outputs["metrics"].items():
+                                if isinstance(value, torch.Tensor):
+                                    task_metrics[task_name][metric_name].append(value.cpu().item())
+                                elif isinstance(value, dict):
+                                    for sub_metric, sub_value in value.items():
+                                        combined_key = f"{metric_name}_{sub_metric}"
+                                        if isinstance(sub_value, (int, float, torch.Tensor)):
+                                            if isinstance(sub_value, torch.Tensor):
+                                                task_metrics[task_name][combined_key].append(
+                                                    sub_value.item()
+                                                )
+                                            else:
+                                                task_metrics[task_name][combined_key].append(
+                                                    sub_value
+                                                )
+                                else:
+                                    task_metrics[task_name][metric_name].append(value)
+
+        if val_batches > 0:
+            avg_loss = total_loss / val_batches
+            aggregated_metrics = {}
+            for task, metrics in task_metrics.items():
+                aggregated_metrics[task] = {}
+                for metric_name, values in metrics.items():
+                    if values:
+                        aggregated_metrics[task][metric_name] = sum(values) / len(values)
+            return {"avg_loss": avg_loss, "metrics": aggregated_metrics}
+        return None

src/dytr/training/utils.py

@@ -0,0 +1,196 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Training utilities for Dynamic Transformers.
+
+This module provides helper functions for training, including seed setting
+and dataset processing for QA tasks.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+
+
+
+import random
+from typing import Dict, List, Optional
+
+import numpy as np
+#import pandas as pd
+import torch
+
+
+def set_seed(seed: int):
+    """
+    Set random seeds for reproducibility.
+
+    Args:
+        seed: Random seed value
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def process_qa_dataset(
+    ds,
+    config,
+    num_samples: int = 200000,
+    text_column: Optional[str] = None,
+    conversations_col: Optional[str] = None,
+    question_col: Optional[str] = None,
+    answer_col: Optional[str] = None,
+    min_text_length: int = 50, 
+) -> List[Dict[str, str]]:
+    """
+    Process question-answering datasets into text format for language modeling.
+
+    This function processes QA datasets from various formats (conversations,
+    question-answer pairs) into a unified text format suitable for training.
+
+    Args:
+        ds: Dataset (list of dictionaries or DataFrame)
+        config: Model configuration (for special tokens)
+        num_samples: Maximum number of samples to process
+        text_column: Column containing text (if applicable)
+        conversations_col: Column containing conversation history
+        question_col: Column containing questions
+        answer_col: Column containing answers
+        min_text_length: Minimum text length to include
+
+    Returns:
+        List of processed text samples
+    """
+    processed_texts = []
+    all_processed_texts = []
+    answer_start = config.special_tokens.get("answer_start", "<|answer|>")
+
+    c = 0
+    for i, item in enumerate(ds):
+        if c >= 20000:
+            c = 0
+            all_processed_texts.extend([{"text": t.lower()} for t in processed_texts])
+            processed_texts = []
+
+        try:
+            # Handle conversation format
+            if conversations_col in item and isinstance(item[conversations_col], list):
+                conversations = item[conversations_col]
+                if len(conversations) >= 2:
+                    human_text = ""
+                    gpt_text = ""
+                    for conv in conversations:
+                        if isinstance(conv, dict):
+                            from_val = conv.get("from", "").lower()
+                            value_val = conv.get("value", "")
+                            if "human" in from_val or "user" in from_val or "question" in from_val:
+                                human_text = value_val
+                            elif (
+                                "gpt" in from_val or "assistant" in from_val or "answer" in from_val
+                            ):
+                                gpt_text = value_val
+
+                    if human_text and gpt_text:
+                        merged = f"Q: {human_text} \n {answer_start} \n {gpt_text} \n \n"
+                        if len(merged) >= min_text_length:
+                            processed_texts.append(merged)
+                            c += len(merged)
+
+            # Handle question-answer format
+            elif (
+                question_col
+                and answer_col
+                and question_col in item
+                and item[question_col]
+                and answer_col in item
+                and item[answer_col]
+            ):
+                question = item[question_col]
+                if text_column and text_column in item and item[text_column]:
+                    question = question + f"\n {item[text_column]}"
+
+                answer = item[answer_col]
+                if question and answer and isinstance(question, str) and isinstance(answer, str):
+                    merged = f"Q: {question} \n {answer_start} \n {answer} \n \n"
+                    if len(merged) >= min_text_length:
+                        processed_texts.append(merged)
+                        c += len(merged)
+
+            # Try to auto-detect question and answer columns
+            else:
+                question = None
+                answer = None
+                question_candidates = [
+                    "question",
+                    "Question",
+                    "query",
+                    "Query",
+                    "input",
+                    "Input",
+                    "prompt",
+                    "Prompt",
+                    "instruction",
+                    "Instruction",
+                ]
+                answer_candidates = [
+                    "answer",
+                    "Answer",
+                    "response",
+                    "Response",
+                    "output",
+                    "Output",
+                    "completion",
+                    "Completion",
+                    "target",
+                    "Target",
+                ]
+
+                for q_col in question_candidates:
+                    if q_col in item and isinstance(item[q_col], str) and item[q_col].strip():
+                        question = item[q_col]
+                        break
+
+                for a_col in answer_candidates:
+                    if a_col in item and isinstance(item[a_col], str) and item[a_col].strip():
+                        answer = item[a_col]
+                        break
+
+                if not question and not answer and text_column in item and item[text_column]:
+                    text = item[text_column]
+                    if isinstance(text, str) and len(text.strip()) >= min_text_length:
+                        processed_texts.append(text + "\n\n")
+                        c += len(text)
+
+                elif question and answer:
+                    merged = f"Q: {question} \n {answer_start} \n {answer} \n \n"
+                    if len(merged) >= min_text_length:
+                        processed_texts.append(merged)
+                        c += len(merged)
+
+            if i == 0 and processed_texts:
+                print(f"  Sample processed text: {processed_texts[0]}...")
+
+        except Exception as e:
+            print(f"  Error in QA processing: {e}")
+            continue
+
+    if processed_texts:
+        all_processed_texts.extend([{"text": t.lower()} for t in processed_texts])
+        #all_processed_texts.append({"text": "\n".join(processed_texts)})
+        print(f"Total Questions Answers: {len(all_processed_texts)}")
+
+    return all_processed_texts

src/dytr/utils/__init__.py

@@ -0,0 +1,31 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Utility modules for Dynamic Transformers.
+
+This module provides utility functions and classes for caching, file management,
+and other helper functionalities.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+__version__ = "0.1.0"
+__author__ = "Akram Alsubari"
+__email__ = "akram.alsubari@outlook.com"
+__license__ = "Apache 2.0"
+
+from dytr.utils.caching import CacheManager, get_cache_manager
+from dytr.utils.logging import get_logger, set_log_level,disable_logging
+__all__ = ["CacheManager", "get_cache_manager","get_logger","set_log_level","disable_logging"]

src/dytr/utils/caching.py

@@ -0,0 +1,101 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Caching utilities for dataset and model caching.
+
+This module provides helper functions for managing cache directories
+and file operations.
+
+Author: Akram Alsubari
+Email: akram.alsubari@outlook.com / akram.alsubari87@gmail.com
+"""
+__version__ = "0.1.0"
+__author__ = "Akram Alsubari"
+__email__ = "akram.alsubari@outlook.com"
+__license__ = "Apache 2.0"
+#import os
+from pathlib import Path
+from typing import Optional
+
+
+class CacheManager:
+    """
+    Manages cache directories for datasets and models.
+
+    Args:
+        cache_dir: Root cache directory path
+    """
+
+    def __init__(self, cache_dir: str = "./dataset_cache"):
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def get_dataset_cache_path(
+        self, dataset_name: str, task_name: str, max_len: int, stride: int
+    ) -> Path:
+        """Get cache path for a processed dataset."""
+        cache_subdir = self.cache_dir / "datasets" / task_name
+        cache_subdir.mkdir(parents=True, exist_ok=True)
+        filename = f"{dataset_name}_{max_len}_{stride}.pkl"
+        return cache_subdir / filename
+
+    def get_model_cache_path(self, model_name: str, task_name: str) -> Path:
+        """Get cache path for a model."""
+        cache_subdir = self.cache_dir / "models" / task_name
+        cache_subdir.mkdir(parents=True, exist_ok=True)
+        return cache_subdir / f"{model_name}.pt"
+
+    def get_tokenizer_cache_path(self, tokenizer_name: str) -> Path:
+        """Get cache path for a tokenizer."""
+        cache_subdir = self.cache_dir / "tokenizers"
+        cache_subdir.mkdir(parents=True, exist_ok=True)
+        return cache_subdir / tokenizer_name
+
+    def clear_cache(self, cache_type: Optional[str] = None):
+        """Clear cache files."""
+        if cache_type == "datasets":
+            cache_path = self.cache_dir / "datasets"
+        elif cache_type == "models":
+            cache_path = self.cache_dir / "models"
+        elif cache_type == "tokenizers":
+            cache_path = self.cache_dir / "tokenizers"
+        else:
+            cache_path = self.cache_dir
+
+        if cache_path.exists():
+            import shutil
+
+            shutil.rmtree(cache_path)
+            print(f"Cleared cache at {cache_path}")
+
+    def get_cache_size(self) -> int:
+        """Get total cache size in bytes."""
+        total_size = 0
+        for file_path in self.cache_dir.rglob("*"):
+            if file_path.is_file():
+                total_size += file_path.stat().st_size
+        return total_size
+
+
+# Global cache manager instance
+_default_cache_manager = None
+
+
+def get_cache_manager(cache_dir: str = "./dataset_cache") -> CacheManager:
+    """Get or create a cache manager instance."""
+    global _default_cache_manager
+    if _default_cache_manager is None or _default_cache_manager.cache_dir != Path(cache_dir):
+        _default_cache_manager = CacheManager(cache_dir)
+    return _default_cache_manager

src/dytr/utils/logging.py

@@ -0,0 +1,64 @@
+# Copyright 2025 Akram Alsubari
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+
+import logging
+import sys
+from typing import Optional
+
+_DEFAULT_LOG_LEVEL = logging.INFO
+_ROOT_LOGGER_NAME = "dytr"
+
+def get_logger(name: str, level: Optional[int] = None) -> logging.Logger:
+    """
+    Get a logger for a module with consistent formatting.
+    
+    Args:
+        name: Logger name (typically __name__)
+        level: Logging level (default: INFO)
+        
+    Returns:
+        Configured logger instance
+    """
+    logger = logging.getLogger(name)
+    
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stderr)
+        formatter = logging.Formatter(
+            '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S'
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+        
+        if level is None:
+            level = _DEFAULT_LOG_LEVEL
+        logger.setLevel(level)
+        logger.propagate = False
+    
+    return logger
+
+def set_log_level(level: int):
+    """
+    Set log level for all dytr loggers.
+    
+    Args:
+        level: Logging level (logging.DEBUG, logging.INFO, etc.)
+    """
+    logging.getLogger(_ROOT_LOGGER_NAME).setLevel(level)
+
+def disable_logging():
+    """Disable all dytr logging."""
+    logging.getLogger(_ROOT_LOGGER_NAME).setLevel(logging.WARNING)