Views
No views yet
1from transformers import AutoTokenizer
2from auto_gptq import AutoGPTQForCausalLM
3
4
5def generate_answer(model, tokenizer, request: str, system_prompt: str):
6 s = f"system\n{system_prompt}</s>\n" + \
7 f"<s>user\n{request}</s>\n" + \
8 f"<s>bot\n"
9 request_tokens = tokenizer(s, return_tensors="pt")
10 del request_tokens['token_type_ids']
11 del request_tokens['attention_mask']
12 request_tokens = request_tokens.to(model.device)
13 answer_tokens = model.generate(**request_tokens,
14 num_beams=4,
15 top_k=32,
16 temperature=0.6,
17 repetition_penalty=1.2,
18 no_repeat_ngram_size=15,
19 max_new_tokens=1536,
20 pad_token_id=tokenizer.eos_token_id)[0]
21 answer_tokens = answer_tokens[len(request_tokens[0]):-1]
22 answer = tokenizer.decode(answer_tokens).strip()
23 return answer
24
25
26model_name = "saiga2-13b-4bit"
27system_prompt = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
28tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
29model = AutoGPTQForCausalLM.from_quantized(model_name, device="cuda:0")
30model.eval()
31
32user_text = "Сочини стих, который начинается словами: Буря мглою небо кроет"
33answer_text = generate_answer(model, tokenizer, user_text, system_prompt)
34print(answer_text)