Upload folder using huggingface_hub

checkpoints/checkpoint_final.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5944f5d325c8c7f6e666aca78f6883eab157f9d416fc4edd748f0308a2c409f0
+size 284829424

infer.py

@@ -0,0 +1,70 @@
+# ==========================
+# infer.py
+# ==========================
+# Usage:
+#   python infer.py --checkpoint checkpoints/checkpoint_epoch3.pt --spm_model spm.model
+
+import argparse
+import torch
+import sentencepiece as spm
+from train import GPT, GPTConfig
+
+def top_k_logits(logits, k):
+    v, ix = torch.topk(logits, k)
+    out = logits.clone()
+    out[out < v[..., -1, None]] = -float('Inf')
+    return out
+
+def generate(model, sp, device, prompt, max_new_tokens=128, temperature=1.0, top_k=50):
+    model.eval()
+    ids = sp.EncodeAsIds('<s>' + prompt + '<sep>')
+    context = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+
+    for _ in range(max_new_tokens):
+        seq = context if context.size(1) <= model.block_size else context[:, -model.block_size:]
+        logits, _ = model(seq)
+        logits = logits[:, -1, :] / max(temperature, 1e-8)
+        if top_k is not None:
+            logits = top_k_logits(logits, top_k)
+        probs = torch.softmax(logits, dim=-1)
+        next_id = torch.multinomial(probs, num_samples=1)
+        context = torch.cat([context, next_id], dim=1)
+        if next_id.item() == sp.EncodeAsIds('</s>')[0]:
+            break
+    out_ids = context[0].tolist()
+    text = sp.DecodeIds([i for i in out_ids if i != 0])
+    if '<sep>' in text:
+        parts = text.split('<sep>')
+        if len(parts) >= 2:
+            return parts[1].replace('</s>','').strip()
+    return text
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint', type=str, required=True)
+    parser.add_argument('--spm_model', type=str, default='spm.model')
+    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
+    parser.add_argument('--prompt', type=str, default='你好，请自我介绍。')
+    parser.add_argument('--max_new_tokens', type=int, default=1024)
+    parser.add_argument('--temperature', type=float, default=0.8)
+    parser.add_argument('--top_k', type=int, default=40)
+    args = parser.parse_args()
+
+    sp = spm.SentencePieceProcessor()
+    sp.Load(args.spm_model)
+
+    ckpt = torch.load(args.checkpoint, map_location=args.device)
+    cfg = ckpt.get('config')
+    config = GPTConfig(vocab_size=cfg['vocab_size'], n_layer=cfg['n_layer'], n_head=cfg['n_head'], n_embd=cfg['n_embd'], block_size=cfg['block_size'], dropout=cfg['dropout'])
+    model = GPT(config).to(args.device)
+    model.load_state_dict(ckpt['model_state'])
+
+    param_count = sum(p.numel() for p in model.parameters())
+    print(f"Loaded model parameters: {param_count:,} ({param_count/1e9:.3f} B)")
+
+    ans = generate(model, sp, args.device, args.prompt, max_new_tokens=args.max_new_tokens, temperature=args.temperature, top_k=args.top_k)
+    print('\n=== RESPONSE ===\n')
+    print(ans)
+
+if __name__ == '__main__':
+    main()

spm.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:478beeaf55eb5dedd0b87f0cd5c032341faec7b3500f7636d55407f5cc6cc809
+size 989403

train.py

@@ -0,0 +1,237 @@
+# ==========================
+# train.py
+# ==========================
+# Usage:
+#   python train.py --data_path all.jsonl --spm_model spm.model
+# Requirements:
+#   pip install torch sentencepiece tqdm
+
+import os
+import json
+import sentencepiece as spm
+from argparse import ArgumentParser
+from tqdm import tqdm
+
+import torch
+from torch import nn
+from torch.utils.data import Dataset, DataLoader
+
+# --------------------------
+# Simple Decoder-only Transformer (GPT-like)
+# --------------------------
+
+class GPTConfig:
+    def __init__(self, vocab_size, n_layer=12, n_head=12, n_embd=768, block_size=1024, dropout=0.1):
+        self.vocab_size = vocab_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.block_size = block_size
+        self.dropout = dropout
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        self.n_head = config.n_head
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x, attn_mask=None):
+        B, T, C = x.size()
+        qkv = self.c_attn(x)  # (B, T, 3*C)
+        q, k, v = qkv.split(C, dim=2)
+        # reshape for multi-head
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1,2)  # (B, nh, T, hs)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1,2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1,2)
+
+        att = (q @ k.transpose(-2, -1)) / (C // self.n_head) ** 0.5  # (B, nh, T, T)
+        # causal mask
+        mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
+        att = att.masked_fill(mask == 0, float('-inf'))
+        att = torch.softmax(att, dim=-1)
+        att = self.dropout(att)
+
+        y = att @ v  # (B, nh, T, hs)
+        y = y.transpose(1,2).contiguous().view(B, T, C)
+        y = self.c_proj(y)
+        y = self.dropout(y)
+        return y
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = nn.LayerNorm(config.n_embd)
+        self.mlp = nn.Sequential(
+            nn.Linear(config.n_embd, 4 * config.n_embd),
+            nn.GELU(),
+            nn.Linear(4 * config.n_embd, config.n_embd),
+            nn.Dropout(config.dropout),
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.block_size = config.block_size
+
+        # initialize
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.size()
+        assert T <= self.block_size
+        token_embeddings = self.tok_emb(idx)  # (B, T, C)
+        x = token_embeddings + self.pos_emb[:, :T, :]
+        x = self.drop(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+        logits = self.head(x)
+
+        loss = None
+        if targets is not None:
+            # shift logits and targets for next-token prediction
+            logits = logits[:, :-1, :].contiguous()
+            targets = targets[:, 1:].contiguous()
+            loss = nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+
+# --------------------------
+# Dataset and helpers
+# --------------------------
+
+class QADataset(Dataset):
+    def __init__(self, path, sp_model, block_size=1024):
+        self.examples = []
+        self.block_size = block_size
+        self.sp = sp_model
+        with open(path, 'r', encoding='utf-8') as f:
+            for line in f:
+                obj = json.loads(line)
+                q = obj.get('question','')
+                a = obj.get('answer','')
+                # format: <bos> question <sep> answer <eos>
+                text = "<s>" + q + "<sep>" + a + "</s>"
+                ids = self.sp.EncodeAsIds(text)
+                if len(ids) > 2:
+                    # truncate or pad later
+                    self.examples.append(ids)
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        ids = self.examples[idx]
+        # pad/truncate to block_size
+        if len(ids) > self.block_size:
+            ids = ids[:self.block_size]
+        else:
+            ids = ids + [0] * (self.block_size - len(ids))
+        return torch.tensor(ids, dtype=torch.long)
+
+
+def collate_fn(batch):
+    batch = torch.stack(batch, dim=0)
+    return batch, batch  # inputs and targets are same sequence for causal LM
+
+
+# --------------------------
+# Main training loop
+# --------------------------
+
+
+def train(args):
+    # prepare sentencepiece model (if not exists, train it)
+    if not os.path.exists(args.spm_model):
+        print('Training SentencePiece model...')
+        # create a temporary file with concatenated text
+        tmp_txt = 'spm_input.txt'
+        with open(args.data_path, 'r', encoding='utf-8') as fin, open(tmp_txt, 'w', encoding='utf-8') as fout:
+            for line in fin:
+                obj = json.loads(line)
+                text = obj.get('question','') + '\n' + obj.get('answer','') + '\n'
+                fout.write(text)
+        spm.SentencePieceTrainer.Train(f'--input={tmp_txt} --model_prefix=spm --vocab_size={args.vocab_size} --model_type=bpe --character_coverage=0.9995')
+        os.remove(tmp_txt)
+        sp = spm.SentencePieceProcessor()
+        sp.Load('spm.model')
+    else:
+        sp = spm.SentencePieceProcessor()
+        sp.Load(args.spm_model)
+
+    dataset = QADataset(args.data_path, sp, block_size=args.block_size)
+    print(f"Loaded {len(dataset)} examples")
+    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, collate_fn=lambda x: collate_fn(x))
+
+    config = GPTConfig(vocab_size=args.vocab_size, n_layer=args.n_layer, n_head=args.n_head, n_embd=args.n_embd, block_size=args.block_size, dropout=args.dropout)
+    model = GPT(config).to(args.device)
+
+    # print parameter count
+    param_count = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {param_count:,} ({param_count/1e9:.3f} B)")
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+
+    model.train()
+    for epoch in range(args.epochs):
+        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{args.epochs}")
+        for batch_inputs, batch_targets in pbar:
+            batch_inputs = batch_inputs.to(args.device)
+            batch_targets = batch_targets.to(args.device)
+            logits, loss = model(batch_inputs, targets=batch_targets)
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            pbar.set_postfix(loss=loss.item())
+
+        # save checkpoint each epoch
+        os.makedirs(args.out_dir, exist_ok=True)
+        torch.save({'model_state': model.state_dict(), 'sp_model': args.spm_model, 'config': vars(config)}, os.path.join(args.out_dir, f'checkpoint_final.pt'))
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument('--data_path', type=str, default='all.jsonl')
+    parser.add_argument('--spm_model', type=str, default='spm.model')
+    parser.add_argument('--vocab_size', type=int, default=32000)
+    parser.add_argument('--block_size', type=int, default=1024)
+    parser.add_argument('--n_layer', type=int, default=3)
+    parser.add_argument('--n_head', type=int, default=3)
+    parser.add_argument('--n_embd', type=int, default=768)
+    parser.add_argument('--batch_size', type=int, default=30)
+    parser.add_argument('--epochs', type=int, default=300)
+    parser.add_argument('--lr', type=float, default=3e-4)
+    parser.add_argument('--dropout', type=float, default=0.1)
+    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
+    parser.add_argument('--out_dir', type=str, default='checkpoints')
+    args = parser.parse_args()
+    train(args)
+