|
--- |
|
language: |
|
- ru |
|
- en |
|
--- |
|
```python |
|
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig |
|
|
|
with ('generation_config.json').open('w') as fp: |
|
json.dump({ |
|
"pad_token_id": 0, |
|
"bos_token_id": 1, |
|
"eos_token_id": 2, |
|
"temperature": 0.3, |
|
"top_p": 0.9, |
|
"top_k": 50, |
|
"do_sample": True, |
|
"max_new_tokens": 1536, |
|
"repetition_penalty": 1.1, |
|
"no_repeat_ngram_size": 15, |
|
}, fp, indent=4) |
|
|
|
MODEL_NAME = "Vikhrmodels/Vikhr_instruct" |
|
TEMPLATE = "<s>{role}\n{content}</s>\n" |
|
SYSTEM_PROMPT = "Ты – полезный помощник по имени Вихрь. Ты разговариваешь с людьми и помогаешь им." |
|
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) |
|
model.to('cuda') |
|
model.eval() |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False) |
|
generation_config = GenerationConfig.from_pretrained("generation_config.json") |
|
|
|
class Conversation: |
|
def __init__( |
|
self, |
|
message_template=DEFAULT_MESSAGE_TEMPLATE, |
|
system_prompt=DEFAULT_SYSTEM_PROMPT, |
|
): |
|
self.message_template = message_template |
|
self.messages = [{ |
|
"role": "system", |
|
"content": system_prompt |
|
}] |
|
|
|
def add_user_message(self, message): |
|
self.messages.append({ |
|
"role": "user", |
|
"content": message |
|
}) |
|
|
|
def get_prompt(self, tokenizer): |
|
final_text = "" |
|
for message in self.messages: |
|
message_text = self.message_template.format(**message) |
|
final_text += message_text |
|
final_text += 'bot' |
|
return final_text.strip() |
|
|
|
|
|
def generate(model, tokenizer, prompt, generation_config): |
|
data = tokenizer(prompt, return_tensors="pt") |
|
data = {k: v.to(model.device) for k, v in data.items()} |
|
output_ids = model.generate( |
|
**data, |
|
generation_config=generation_config |
|
)[0] |
|
output_ids = output_ids[len(data["input_ids"][0]):] |
|
output = tokenizer.decode(output_ids, skip_special_tokens=True) |
|
return output.strip() |
|
|
|
inputs = ["Как тебя зовут?", "Кто такой Колмогоров?"] |
|
|
|
for inp in inputs: |
|
conversation = Conversation() |
|
conversation.add_user_message(inp) |
|
prompt = conversation.get_prompt(tokenizer) |
|
|
|
output = generate(model, tokenizer, prompt, generation_config) |
|
print(inp) |
|
print(output) |
|
``` |
|
|
|
[wandb](https://wandb.ai/karina_romanova/vikhr/runs/up2hw5eh?workspace=user-karina_romanova) |
