Create user.py

user.py

@@ -0,0 +1,220 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json, os, numpy as np
+
+# ============================================================================
+class ChunkTokenizer:
+    def __init__(self):
+        self.chunk_to_idx = {}
+        self.idx_to_chunk = {}
+        self.vocab_size = 0
+
+    def load(self, path):
+        with open(path, 'r') as f:
+            vocab_data = json.load(f)
+        self.chunk_to_idx = vocab_data['chunk_to_idx']
+        self.idx_to_chunk = {int(k): v for k, v in vocab_data['idx_to_chunk'].items()}
+        self.vocab_size = vocab_data['vocab_size']
+        print(f"Loaded tokenizer ({self.vocab_size} tokens)")
+
+    def encode(self, text):
+        text = text.lower()
+        pos, indices = 0, []
+        while pos < len(text):
+            for size in (3, 2, 1):
+                chunk = text[pos:pos+size]
+                if chunk in self.chunk_to_idx:
+                    indices.append(self.chunk_to_idx[chunk])
+                    pos += size
+                    break
+            else:
+                pos += 1
+        return indices
+
+    def decode(self, indices):
+        return ''.join([self.idx_to_chunk.get(int(i), '') for i in indices])
+
+
+# ============================================================================
+class LoRPtLinear(nn.Module):
+    def __init__(self, in_features, out_features, rank=64):
+        super().__init__()
+        self.lora_A = nn.Parameter(torch.randn(out_features, rank) * 0.02)
+        self.lora_B = nn.Parameter(torch.randn(rank, in_features) * 0.02)
+        self.bias = nn.Parameter(torch.zeros(out_features))
+
+    def forward(self, x):
+        return F.linear(x, self.lora_A @ self.lora_B, self.bias)
+
+
+class RWKVMambaHybrid(nn.Module):
+    def __init__(self, d_model, d_state=32):
+        super().__init__()
+        self.d_model = d_model
+        self.d_state = d_state
+        self.w_mix = nn.Parameter(torch.ones(d_model) * 0.5)
+        self.A = nn.Parameter(torch.randn(d_state, d_state) * 0.01)
+        self.B = nn.Parameter(torch.randn(d_state, d_model) * 0.01)
+        self.C = nn.Parameter(torch.randn(d_model, d_state) * 0.01)
+        self.D = nn.Parameter(torch.ones(d_model) * 0.1)
+
+    def forward(self, x):
+        B, T, C = x.shape
+        h = torch.zeros(B, C, device=x.device)
+        s = torch.zeros(B, self.d_state, device=x.device)
+        outputs = []
+        for t in range(T):
+            x_t = x[:, t, :]
+            h = self.w_mix * h + (1 - self.w_mix) * x_t
+            s = s @ self.A.T + x_t @ self.B.T
+            y_t = s @ self.C.T + h * self.D
+            outputs.append(y_t)
+        return torch.stack(outputs, dim=1)
+
+
+class KQVAttention(nn.Module):
+    def __init__(self, d_model, n_heads=16, rank=64):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.q_down = nn.Linear(d_model, rank)
+        self.q_up = nn.Linear(rank, d_model)
+        self.k_down = nn.Linear(d_model, rank)
+        self.k_up = nn.Linear(rank, d_model)
+        self.v_down = nn.Linear(d_model, rank)
+        self.v_up = nn.Linear(rank, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape
+        q = self.q_up(self.q_down(x))
+        k = self.k_up(self.k_down(x))
+        v = self.v_up(self.v_down(x))
+        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        attn = (q @ k.transpose(-2, -1)) / np.sqrt(self.head_dim)
+        if mask is not None:
+            attn = attn.masked_fill(mask == 0, float('-inf'))
+        attn = F.softmax(attn, dim=-1)
+        out = attn @ v
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(out)
+
+
+class i3Block(nn.Module):
+    def __init__(self, d_model, n_heads=16, d_state=32, rank=64, ffn_mult=4):
+        super().__init__()
+        self.hybrid = RWKVMambaHybrid(d_model, d_state)
+        self.ln1 = nn.LayerNorm(d_model)
+        self.attn = KQVAttention(d_model, n_heads, rank)
+        self.ln2 = nn.LayerNorm(d_model)
+        d_ff = d_model * ffn_mult
+        self.ffn = nn.Sequential(
+            LoRPtLinear(d_model, d_ff, rank),
+            nn.GELU(),
+            LoRPtLinear(d_ff, d_model, rank)
+        )
+        self.ln3 = nn.LayerNorm(d_model)
+
+    def forward(self, x, mask=None):
+        x = x + self.hybrid(self.ln1(x))
+        x = x + self.attn(self.ln2(x), mask)
+        x = x + self.ffn(self.ln3(x))
+        return x
+
+
+class i3Model(nn.Module):
+    def __init__(self, vocab_size, d_model=512, n_layers=24, n_heads=16,
+                 max_seq_len=256, rank=64, d_state=32):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+        self.embed = nn.Embedding(vocab_size, d_model)
+        self.pos_embed = nn.Embedding(max_seq_len, d_model)
+        self.layers = nn.ModuleList([
+            i3Block(d_model, n_heads, d_state, rank)
+            for _ in range(n_layers)
+        ])
+        self.ln_f = nn.LayerNorm(d_model)
+        self.head = LoRPtLinear(d_model, vocab_size, rank)
+
+    def forward(self, idx):
+        B, T = idx.shape
+        pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
+        x = self.embed(idx) + self.pos_embed(pos)
+        mask = torch.tril(torch.ones(T, T, device=idx.device)).view(1, 1, T, T)
+        for layer in self.layers:
+            x = layer(x, mask)
+        x = self.ln_f(x)
+        return self.head(x)
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=100, temperature=0.8, top_k=40):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.max_seq_len:]
+            logits = self(idx_cond)[:, -1, :] / temperature
+            v, _ = torch.topk(logits, top_k)
+            logits[logits < v[:, [-1]]] = -float("inf")
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, 1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx
+
+
+# ============================================================================
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+tokenizer = ChunkTokenizer()
+tokenizer.load("tokenizer.json")
+
+model = i3Model(
+    vocab_size=tokenizer.vocab_size,
+    d_model=512,
+    n_layers=24,
+    n_heads=16,
+    max_seq_len=256,
+    rank=64,
+    d_state=32
+).to(device)
+
+state_dict = torch.load("pytorch_model.bin", map_location=device)
+model.load_state_dict(state_dict)
+model.eval()
+print("✓ Model loaded successfully")
+
+
+# ============================================================================
+@torch.no_grad()
+def infer(prompt, max_new_tokens=100, temperature=0.8, top_k=40):
+    input_ids = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long).to(device)
+    output = model.generate(input_ids, max_new_tokens=max_new_tokens,
+                            temperature=temperature, top_k=top_k)
+    return tokenizer.decode(output[0].cpu().numpy())
+
+
+def chat_loop():
+    print("=== i3 Interactive Chat ([INST] format) ===")
+    history = ""
+    while True:
+        user_input = input("[You] ")
+        if user_input.strip().lower() in {"quit", "exit"}:
+            break
+        prompt = f"{history}[INST] {user_input.strip()} [/INST]"
+        reply = infer(prompt, max_new_tokens=120)
+        reply_clean = reply.replace(prompt.lower(), "").strip()
+        print("[i3]:", reply_clean)
+        history += f"[INST] {user_input.strip()} [/INST] {reply_clean} "
+
+
+# ============================================================================
+print("\nExample:")
+prompt = "[INST] What can we do to make people happier [/INST]"
+print("Prompt:", prompt)
+print("Generated:", infer(prompt))
+
+# Optionally start a chat loop:
+# chat_loop()