Upload 3 files

model_adapters.py

@@ -1,5 +1,6 @@
 import json
 import os
+import glob
 from typing import Any, Dict, List
 
 
@@ -105,12 +106,126 @@ class TransformersAdapter:
 
 class NanochatAdapter:
     def __init__(self, model_id: str):
-        raise RuntimeError(
-            "RFAB_HISTORIC_ADAPTER=nanochat is a scaffold. Copy the nanochat runtime "
-            "from tventurella/mr_chatterbox into this Space and replace NanochatAdapter "
-            "with that loader before setting DRY_RUN=false."
+        self.model_id = model_id
+        self.model = None
+        self.tokenizer = None
+        self.device = None
+
+    def _ensure_loaded(self):
+        if self.model is not None and self.tokenizer is not None:
+            return
+
+        import torch
+        from huggingface_hub import snapshot_download
+        from nanochat.gpt import GPT, GPTConfig
+        from nanochat.tokenizer import RustBPETokenizer
+
+        local_dir = snapshot_download(
+            repo_id=self.model_id,
+            token=os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN"),
         )
 
+        tokenizer_dir = os.getenv("RFAB_NANOCHAT_TOKENIZER_DIR", "tokenizer")
+        tokenizer_path = os.path.join(local_dir, tokenizer_dir)
+        self.tokenizer = RustBPETokenizer.from_directory(tokenizer_path)
+
+        meta_file = os.getenv("RFAB_NANOCHAT_META_FILE") or first_match(local_dir, "meta_*.json")
+        model_file = os.getenv("RFAB_NANOCHAT_MODEL_FILE") or first_match(local_dir, "model_*.pt")
+        meta_path = file_in_snapshot(local_dir, meta_file)
+        model_path = file_in_snapshot(local_dir, model_file)
+
+        with open(meta_path, "r", encoding="utf-8") as f:
+            meta = json.load(f)
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        config = GPTConfig(**meta["model_config"])
+        with torch.device("meta"):
+            model = GPT(config)
+
+        model.to_empty(device=self.device)
+        model.init_weights()
+        state_dict = torch.load(model_path, map_location=self.device)
+        state_dict = {k.removeprefix("_orig_mod."): v for k, v in state_dict.items()}
+        model.load_state_dict(state_dict, strict=True, assign=True)
+        model.eval()
+        self.model = model
+
+    def generate(self, history, system_prompt, temperature, max_tokens, top_p, top_k):
+        import torch
+
+        self._ensure_loaded()
+
+        bos = self.tokenizer.get_bos_token_id()
+        user_start = self.tokenizer.encode_special("<|user_start|>")
+        user_end = self.tokenizer.encode_special("<|user_end|>")
+        assistant_start = self.tokenizer.encode_special("<|assistant_start|>")
+        assistant_end = self.tokenizer.encode_special("<|assistant_end|>")
+
+        tokens = [bos]
+        if system_prompt:
+            tokens += [user_start]
+            tokens += self.tokenizer.encode(system_prompt.strip())
+            tokens += [user_end, assistant_start, assistant_end]
+
+        for message in history or []:
+            role = message.get("role")
+            text = extract_text(message)
+            if not text:
+                continue
+            if role == "user":
+                tokens += [user_start]
+                tokens += self.tokenizer.encode(text)
+                tokens += [user_end]
+            elif role == "assistant":
+                tokens += [assistant_start]
+                tokens += self.tokenizer.encode(text)
+                tokens += [assistant_end]
+
+        tokens += [assistant_start]
+        generated = []
+
+        generate_kwargs = {
+            "max_tokens": int(max_tokens),
+            "temperature": float(temperature),
+        }
+        if int(top_k) > 0:
+            generate_kwargs["top_k"] = int(top_k)
+
+        if self.device == "cuda":
+            context = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
+        else:
+            context = nullcontext()
+
+        with torch.no_grad(), context:
+            for token in self.model.generate(tokens, **generate_kwargs):
+                if token in (assistant_end, bos):
+                    break
+                generated.append(token)
+
+        return self.tokenizer.decode(generated).strip()
+
+
+class nullcontext:
+    def __enter__(self):
+        return None
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+
+def first_match(root, pattern):
+    matches = sorted(glob.glob(os.path.join(root, pattern)))
+    if not matches:
+        raise FileNotFoundError(f"No file matching {pattern} in {root}")
+    return os.path.basename(matches[0])
+
+
+def file_in_snapshot(root, filename):
+    path = filename if os.path.isabs(filename) else os.path.join(root, filename)
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"File not found in snapshot: {filename}")
+    return path
+
 
 def parse_json_env(name: str, default):
     raw = os.getenv(name)

requirements.txt

@@ -5,3 +5,6 @@ torch
 sentencepiece
 huggingface_hub
 spaces
+tiktoken
+rustbpe
+nanochat @ git+https://github.com/karpathy/nanochat.git