1import os
2import torch
3import json
4import tiktoken
5import importlib.util
6from huggingface_hub import hf_hub_download
7
8# === CONFIG ===
9REPO_ID = "faizack/bayes_mini_custom"
10
11# === Step 1: Download necessary files ===
12config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
13model_path = hf_hub_download(repo_id=REPO_ID, filename="pytorch_model.bin")
14modeling_path = hf_hub_download(repo_id=REPO_ID, filename="modeling_gpt2_custom.py")
15
16# === Step 2: Dynamically import modeling_gpt2_custom.py ===
17spec = importlib.util.spec_from_file_location("modeling_gpt2_custom", modeling_path)
18mod = importlib.util.module_from_spec(spec)
19spec.loader.exec_module(mod)
20GPTModel = mod.GPTModel # Now you can use GPTModel
21
22# === Step 3: Load config ===
23with open(config_path, "r") as f:
24 config = json.load(f)
25
26model_config = {
27 "vocab_size": config["vocab_size"],
28 "context_length": config["n_positions"],
29 "emb_dim": config["n_embd"],
30 "n_heads": config["n_head"],
31 "n_layers": config["n_layer"],
32 "drop_rate": config["dropout"],
33 "qkv_bias": config["qkv_bias"],
34}
35
36# === Step 4: Load tokenizer ===
37tokenizer = tiktoken.get_encoding("gpt2")
38prompt = "The rise of artificial intelligence"
39input_ids = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
40
41# === Step 5: Load model ===
42model = GPTModel(model_config)
43model.load_state_dict(torch.load(model_path, map_location="cpu"))
44model.eval()
45
46
47# === Step 6: Generate ===
48def generate(model, idx, max_new_tokens=50):
49 for _ in range(max_new_tokens):
50 idx_cond = idx[:, -model_config["context_length"] :]
51 with torch.no_grad():
52 logits = model(idx_cond)
53 logits = logits[:, -1, :]
54 probs = torch.softmax(logits, dim=-1)
55 next_token = torch.multinomial(probs, num_samples=1)
56 idx = torch.cat([idx, next_token], dim=1)
57 return idx
58
59
60output = generate(model, input_ids)
61print(tokenizer.decode(output[0].tolist()))
62