|
|
| 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() |
