Views
No views yet
Note: If you prefer, you can opt to use the ruGPT-3.5 13B fp16 base model.
1import torch
2from peft import PeftModel, PeftConfig
3from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
4
5MODEL_NAME = "evilfreelancer/ruGPT-3.5-13B-lora"
6DEFAULT_MESSAGE_TEMPLATE = "<s>{role}\n{content}</s>\n"
7DEFAULT_SYSTEM_PROMPT = "Ты — ruGPT-3.5, русскоязычный автоматический ассистент на 13 миллиардов параметров. Ты разговариваешь с людьми и помогаешь им."
8
9class Conversation:
10 def __init__(
11 self,
12 message_template=DEFAULT_MESSAGE_TEMPLATE,
13 system_prompt=DEFAULT_SYSTEM_PROMPT,
14 start_token_id=2,
15 bot_token_id=46787
16 ):
17 self.message_template = message_template
18 self.start_token_id = start_token_id
19 self.bot_token_id = bot_token_id
20 self.messages = [{
21 "role": "system",
22 "content": system_prompt
23 }]
24
25 def get_start_token_id(self):
26 return self.start_token_id
27
28 def get_bot_token_id(self):
29 return self.bot_token_id
30
31 def add_user_message(self, message):
32 self.messages.append({
33 "role": "user",
34 "content": message
35 })
36
37 def add_bot_message(self, message):
38 self.messages.append({
39 "role": "bot",
40 "content": message
41 })
42
43 def get_prompt(self, tokenizer):
44 final_text = ""
45 for message in self.messages:
46 message_text = self.message_template.format(**message)
47 final_text += message_text
48 final_text += tokenizer.decode([self.start_token_id, self.bot_token_id])
49 return final_text.strip()
50
51
52def generate(model, tokenizer, prompt, generation_config):
53 data = tokenizer(prompt, return_tensors="pt")
54 data = {k: v.to(model.device) for k, v in data.items()}
55 output_ids = model.generate(
56 **data,
57 generation_config=generation_config
58 )[0]
59 output_ids = output_ids[len(data["input_ids"][0]):]
60 output = tokenizer.decode(output_ids, skip_special_tokens=True)
61 return output.strip()
62
63config = PeftConfig.from_pretrained(MODEL_NAME)
64model = AutoModelForCausalLM.from_pretrained(
65 config.base_model_name_or_path,
66 load_in_8bit=True,
67 torch_dtype=torch.float16,
68 device_map="auto"
69)
70model = PeftModel.from_pretrained(
71 model,
72 MODEL_NAME,
73 torch_dtype=torch.float16
74)
75model.eval()
76
77tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
78generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
79print(generation_config)
80
81conversation = Conversation()
82while True:
83 user_message = input("User: ")
84
85 if user_message.strip() == "/reset":
86 conversation = Conversation()
87 print("History reset completed!")
88 continue
89
90 if user_message.strip() == "":
91 continue
92
93 conversation.add_user_message(user_message)
94 prompt = conversation.get_prompt(tokenizer)
95 output = generate(
96 model=model,
97 tokenizer=tokenizer,
98 prompt=prompt,
99 generation_config=generation_config
100 )
101 conversation.add_bot_message(output)
102 print("ruGPT-3.5:", output)
103 print()
104 print("==============================")
105 print()1User: Привет! Почему трава зелёная?
2ruGPT-3.5: Трава зелёная благодаря хлорофиллу - пигменту растений, который поглощает световые волны определенной длины и преобразует их в энергию химических связей. Хлорофилл также обеспечивает растения питательными веществами из почвы через фотосинтез.1User: Придумай стишок про нейросети.
2ruGPT-3.5: Нейросеть – это чудо науки,
3Она способна решать любые задачи.
4И если ты хочешь стать программистом,
5То лучше всего изучать нейронные сети.1User: Высота самой высокой горы на Земле?
2ruGPT-3.5: Самая высокая гора на Земле находится в Гималаях и называется Эверест (8848 метров).1User: Сколько максимум лет может прожить галопогосская черепаха?
2ruGPT-3.5: Галопогосские черепахи могут жить до 100 лет и более. Они являются одними из самых долгоживущих животных на планете.bitsandbytes quantization config was used during training: