feat: Add Chinese i18n support

.gitignore

@@ -0,0 +1,36 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+ENV/
+env/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# OS
+.DS_Store
+Thumbs.db

README.md

@@ -75,11 +75,12 @@ The Engram module augments the backbone by retrieving static $N$-gram memory and
   <img width="80%" src="figures/case.png" alt="Long Context Results">
 </p>
 
+
 ## 5. Quick Start
 
 We recommend using Python 3.8+ and PyTorch.
 ```bash
-pip install torch numpy transformers sympy
+pip install torch numpy transformers sympy polib
 ```
 We provide a standalone implementation to demonstrate the core logic of the Engram module:
 ```bash
@@ -88,10 +89,51 @@ python engram_demo_v1.py
 
 > ⚠️ **Note:** The provided code is a demonstration version intended to illustrate the data flow. It mocks standard components (like Attention/MoE/mHC) to focus on the Engram module. 
 
+## 6. Internationalization (i18n) Support 🇬🇧🇨🇳
+
+Engram now supports multiple languages! You can easily switch between English and Chinese (Simplified):
+
+### Usage Example
+
+```python
+from engram import set_language, _
+
+# Set language to Chinese
+set_language('zh_CN')
+
+# Now all user-facing strings will be in Chinese
+print(_('Forward Complete!'))  # Output: 前向传播完成！
+```
+
+### Supported Languages
+
+| Language | Code | Status |
+|----------|------|--------|
+| English | `en` | ✅ Default |
+| 中文 (Chinese Simplified) | `zh_CN` | ✅ Supported |
+
+### Running the Demo
+
+```bash
+# Run the i18n demo
+python examples/i18n_demo.py
+
+# Run tests
+python tests/test_i18n.py
+```
+
+### Adding New Translations
+
+To add support for a new language:
+
+1. Create a new directory under `engram/locales/` (e.g., `engram/locales/fr/LC_MESSAGES/`)
+2. Copy `engram.po` from an existing locale and translate the strings
+3. Compile to `.mo` format using: `msgfmt engram.po -o engram.mo`
+4. Submit a PR!
 
-## 6. License
+## 7. License
 The use of Engram models is subject to [the Model License](LICENSE).
 
-## 7. Contact
+## 8. Contact
 
-If you have any questions, please raise an issue or contact us at [service@deepseek.com](mailto:service@deepseek.com).
+If you have any questions, please raise an issue or contact us at [service@deepseek.com](mailto:service@deepseek.com).

engram/__init__.py

@@ -0,0 +1,20 @@
+"""
+Engram - Conditional Memory via Scalable Lookup
+
+This package provides the Engram module implementation for efficient
+N-gram memory retrieval in transformer models.
+
+Internationalization (i18n) is supported. To use a different language:
+    from engram import set_language
+    set_language('zh_CN')  # 设置为中文
+"""
+
+from engram.i18n import _, set_language, get_current_language, get_locales_dir
+
+__version__ = "0.1.0"
+__all__ = [
+    '_',
+    'set_language',
+    'get_current_language', 
+    'get_locales_dir',
+]

engram/i18n.py

@@ -0,0 +1,76 @@
+"""
+Internationalization (i18n) support for Engram.
+
+This module provides translation capabilities for user-facing strings.
+Usage:
+    from engram.i18n import _, set_language
+
+    # Set language (default: 'en')
+    set_language('zh_CN')
+
+    # Translate a string
+    print(_("Hello, World!"))
+"""
+
+import gettext
+import os
+from typing import Optional
+
+# Default language
+_DEFAULT_LANGUAGE = 'en'
+_CURRENT_LANGUAGE = _DEFAULT_LANGUAGE
+
+# Translation object (initialized lazily)
+_translation = None
+
+def get_locales_dir() -> str:
+    """Get the directory containing locale files."""
+    return os.path.join(os.path.dirname(__file__), 'locales')
+
+def set_language(lang_code: str) -> None:
+    """
+    Set the current language for translations.
+
+    Args:
+        lang_code: Language code (e.g., 'en', 'zh_CN')
+    """
+    global _CURRENT_LANGUAGE, _translation
+
+    _CURRENT_LANGUAGE = lang_code
+
+    locales_dir = get_locales_dir()
+
+    try:
+        _translation = gettext.translation(
+            'engram',
+            localedir=locales_dir,
+            languages=[lang_code],
+            fallback=(lang_code == _DEFAULT_LANGUAGE)
+        )
+    except FileNotFoundError:
+        # Fallback to English if translation file not found
+        _translation = gettext.NullTranslations()
+
+def _(message: str) -> str:
+    """
+    Translate a message.
+
+    Args:
+        message: The message to translate
+
+    Returns:
+        The translated message
+    """
+    global _translation
+
+    if _translation is None:
+        set_language(_DEFAULT_LANGUAGE)
+
+    return _translation.gettext(message)
+
+def get_current_language() -> str:
+    """Get the currently set language code."""
+    return _CURRENT_LANGUAGE
+
+# Initialize with default language on module import
+set_language(_DEFAULT_LANGUAGE)

engram/locales/en/LC_MESSAGES/engram.po

@@ -0,0 +1,48 @@
+# English translations for Engram.
+# Copyright (C) 2024 DeepSeek AI
+# This file is distributed under the same license as the Engram package.
+#
+msgid ""
+msgstr ""
+"Project-Id-Version: Engram 0.1.0\n"
+"Report-Msgid-Bugs-To: service@deepseek.com\n"
+"POT-Creation-Date: 2024-01-01 00:00+0000\n"
+"PO-Revision-Date: 2024-01-01 00:00+0000\n"
+"Last-Translator: DeepSeek AI\n"
+"Language-Team: English\n"
+"Language: en\n"
+"MIME-Version: 1.0\n"
+"Content-Type: text/plain; charset=UTF-8\n"
+"Content-Transfer-Encoding: 8bit\n"
+
+#: engram_demo_v1.py
+msgid "Forward Complete!"
+msgstr "Forward Complete!"
+
+#: engram_demo_v1.py
+msgid "input_ids.shape"
+msgstr "input_ids.shape"
+
+#: engram_demo_v1.py
+msgid "output.shape"
+msgstr "output.shape"
+
+#: engram_demo_v1.py
+msgid "Error: Failed to initialize tokenizer"
+msgstr "Error: Failed to initialize tokenizer"
+
+#: engram_demo_v1.py
+msgid "Error: Invalid input shape"
+msgstr "Error: Invalid input shape"
+
+#: engram_demo_v1.py
+msgid "Engram Demo - Processing text"
+msgstr "Engram Demo - Processing text"
+
+#: engram_demo_v1.py
+msgid "Initializing Engram module"
+msgstr "Initializing Engram module"
+
+#: engram_demo_v1.py
+msgid "Processing complete"
+msgstr "Processing complete"

engram/locales/zh_CN/LC_MESSAGES/engram.po

@@ -0,0 +1,48 @@
+# Chinese (Simplified) translations for Engram.
+# Copyright (C) 2024 DeepSeek AI
+# This file is distributed under the same license as the Engram package.
+#
+msgid ""
+msgstr ""
+"Project-Id-Version: Engram 0.1.0\n"
+"Report-Msgid-Bugs-To: service@deepseek.com\n"
+"POT-Creation-Date: 2024-01-01 00:00+0000\n"
+"PO-Revision-Date: 2024-01-01 00:00+0000\n"
+"Last-Translator: DeepSeek AI\n"
+"Language-Team: Chinese (Simplified)\n"
+"Language: zh_CN\n"
+"MIME-Version: 1.0\n"
+"Content-Type: text/plain; charset=UTF-8\n"
+"Content-Transfer-Encoding: 8bit\n"
+
+#: engram_demo_v1.py
+msgid "Forward Complete!"
+msgstr "前向传播完成！"
+
+#: engram_demo_v1.py
+msgid "input_ids.shape"
+msgstr "输入ID形状"
+
+#: engram_demo_v1.py
+msgid "output.shape"
+msgstr "输出形状"
+
+#: engram_demo_v1.py
+msgid "Error: Failed to initialize tokenizer"
+msgstr "错误：分词器初始化失败"
+
+#: engram_demo_v1.py
+msgid "Error: Invalid input shape"
+msgstr "错误：无效的输入形状"
+
+#: engram_demo_v1.py
+msgid "Engram Demo - Processing text"
+msgstr "Engram 演示 - 处理文本中"
+
+#: engram_demo_v1.py
+msgid "Initializing Engram module"
+msgstr "正在初始化 Engram 模块"
+
+#: engram_demo_v1.py
+msgid "Processing complete"
+msgstr "处理完成"

engram_demo_v1.py

@@ -23,6 +23,9 @@
 """
 
 ## built-in
+import argparse
+import os
+import sys
 from typing import List
 from dataclasses import dataclass, field
 import math
@@ -35,6 +38,22 @@
 from transformers import AutoTokenizer
 from tokenizers import normalizers, Regex 
 
+# Add parent directory to path for i18n import
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+try:
+    from engram.i18n import _, set_language, get_current_language
+    I18N_AVAILABLE = True
+except ImportError:
+    # Fallback if i18n module is not available
+    def _(text):
+        return text
+    def set_language(lang):
+        pass
+    def get_current_language():
+        return 'en'
+    I18N_AVAILABLE = False
+
 @dataclass
 class EngramConfig:
     tokenizer_name_or_path: str = "deepseek-ai/DeepSeek-V3"
@@ -393,14 +412,42 @@ def forward(self,input_ids,hidden_states):
         hidden_states = self.moe(hidden_states) + hidden_states
         return hidden_states
 
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description=_("Engram Architecture Demo with i18n support")
+    )
+    parser.add_argument(
+        '--language', '-l',
+        type=str,
+        default='en',
+        choices=['en', 'zh_CN'],
+        help=_('Language code (en for English, zh_CN for Chinese)')
+    )
+    parser.add_argument(
+        '--text', '-t',
+        type=str,
+        default=None,
+        help=_('Input text for the demo (default: use example sentence)')
+    )
+    return parser.parse_args()
+
 if __name__ == '__main__':
+    args = parse_args()
+
+    # Set language if i18n is available
+    if I18N_AVAILABLE:
+        set_language(args.language)
+        print(f"Language set to: {get_current_language()}")
+
     LLM = [
         nn.Embedding(backbone_config.vocab_size,backbone_config.hidden_size),
         *[TransformerBlock(layer_id=layer_id) for layer_id in range(backbone_config.num_layers)],
         nn.Linear(backbone_config.hidden_size, backbone_config.vocab_size)
     ]
 
-    text = "Only Alexander the Great could tame the horse Bucephalus."
+    # Use provided text or default example
+    text = args.text if args.text else _("Only Alexander the Great could tame the horse Bucephalus.")
     tokenizer = AutoTokenizer.from_pretrained(engram_cfg.tokenizer_name_or_path,trust_remote_code=True)
     input_ids = tokenizer(text,return_tensors='pt').input_ids
 
@@ -418,6 +465,5 @@ def forward(self,input_ids,hidden_states):
         else:
             hidden_states = layer(input_ids=input_ids,hidden_states=hidden_states)
 
-    print("✅ Forward Complete!")
+    print(_("✅ Forward Complete!"))
     print(f"{input_ids.shape=}\n{output.shape=}")
-