File size: 4,508 Bytes
769a891 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | # use.py
import os
import json
import torch
import torch.nn.functional as F
from transformers import GPT2TokenizerFast, GPT2Config
from safetensors.torch import load_file
from model import VDrontModel
CONFIG = {
"model_dir": "./VDrontV3-Mini",
"temperature": 0.4,
"top_k": 50,
"max_new_tokens": 200,
"repetition_penalty": 1.2,
"user_token": "<|user|>",
"assistant_token": "<|assistant|>",
}
def format_prompt(user_input):
return f"{CONFIG['user_token']}{user_input}{CONFIG['assistant_token']}"
def main():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
tokenizer = GPT2TokenizerFast.from_pretrained(CONFIG['model_dir'])
vocab_size = len(tokenizer)
special_tokens = [CONFIG['user_token'], CONFIG['assistant_token']]
tokenizer.add_special_tokens({'additional_special_tokens': special_tokens})
with open(os.path.join(CONFIG['model_dir'], 'architecture.json')) as f:
arch = json.load(f)
config = GPT2Config(
vocab_size=vocab_size,
n_embd=arch['n_embd'],
n_head=arch['n_head'],
n_layer=arch['n_layer'],
n_positions=arch['n_positions'],
layer_norm_epsilon=1e-5,
)
model = VDrontModel(
config=config,
expert_start=arch['expert_start'],
expert_end=arch['expert_end'],
output_index=arch['output_index'],
num_experts=arch['num_experts'],
num_output_versions=arch['num_output_versions'],
)
state = load_file(os.path.join(CONFIG['model_dir'], 'model.safetensors'))
model.load_state_dict(state)
model.to(device)
model.eval()
if model.embed_tokens.num_embeddings < len(tokenizer):
old_embed = model.embed_tokens
new_embed = torch.nn.Embedding(len(tokenizer), old_embed.embedding_dim).to(device)
new_embed.weight.data[:old_embed.num_embeddings] = old_embed.weight.data.to(device)
model.embed_tokens = new_embed
old_lm_head = model.lm_head
new_lm_head = torch.nn.Linear(old_lm_head.in_features, len(tokenizer), bias=False).to(device)
new_lm_head.weight.data[:old_lm_head.out_features] = old_lm_head.weight.data.to(device)
model.lm_head = new_lm_head
model.config.vocab_size = len(tokenizer)
while True:
try:
output_ver = int(input("Mode (0 - base (bad, little answer), 1 - qualitative (normal, medium answer): "))
if output_ver in [0, 1]:
model.set_output_version(output_ver)
break
except ValueError:
pass
print("Chat is ready. Type 'exit' to quit.")
while True:
user_input = input("You: ")
if user_input.lower() in ['exit', 'quit']:
break
prompt = format_prompt(user_input)
input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
generated_tokens = []
eos_id = tokenizer.eos_token_id
with torch.no_grad():
for _ in range(CONFIG['max_new_tokens']):
pos = torch.arange(0, input_ids.size(1), device=device).unsqueeze(0)
x = model.embed_tokens(input_ids) + model.embed_positions(pos)
router_logits = model.router(x.mean(dim=1))
expert_idx = router_logits.argmax(dim=-1).item()
model.set_expert_version(expert_idx)
idx_cond = input_ids[:, -model.config.n_positions:]
logits, _ = model(idx_cond)
logits = logits[:, -1, :] / CONFIG['temperature']
for token_id in set(input_ids[0].tolist()):
logits[0, token_id] /= CONFIG['repetition_penalty']
if CONFIG['top_k'] is not None:
v, _ = torch.topk(logits, min(CONFIG['top_k'], logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
next_token = idx_next.item()
if next_token == eos_id:
break
generated_tokens.append(next_token)
input_ids = torch.cat((input_ids, idx_next), dim=1)
full_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
print(f"AI: {full_text}")
if __name__ == '__main__':
main() |