1import torch
2import re
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5model_path = "ERmak158/Qwen3-4B-eq-finetuned2"
6
7
8def load_model(path: str):
9 print(f"Loading model from: {path}")
10 try:
11 tokenizer = AutoTokenizer.from_pretrained(path)
12 model = AutoModelForCausalLM.from_pretrained(
13 path,
14 torch_dtype=torch.float16 if torch.cuda.is_available() else "auto",
15 device_map="auto",
16 )
17 model.eval()
18 return tokenizer, model
19 except Exception as e:
20 print(f"Load error: {e}")
21 exit(1)
22
23
24def sanitize_message(content):
25 if content is None:
26 return "..."
27 content = str(content).strip()
28 if not content:
29 return "..."
30 return content
31
32
33def sanitize_history(history):
34 clean_history = []
35 for msg in history:
36 if not isinstance(msg, dict):
37 continue
38 role = msg.get("role")
39 content = msg.get("content")
40
41 if role not in ("system", "user", "assistant"):
42 continue
43
44 clean_content = sanitize_message(content)
45 clean_history.append({"role": role, "content": clean_content})
46
47 return clean_history
48
49
50def generate_stream(tokenizer, model, messages, max_new_tokens=1000,
51 temperature=0.4, top_p=0.4, repetition_penalty=1.25):
52
53 safe_messages = sanitize_history(messages)
54
55 if not safe_messages:
56 safe_messages = [{"role": "system", "content": "Ты полезный ассистент."}]
57
58 try:
59 input_ids = tokenizer.apply_chat_template(
60 safe_messages,
61 tokenize=True,
62 add_generation_prompt=True,
63 return_tensors="pt"
64 ).to(model.device)
65 except Exception as e:
66 print(f"Template error: {e}")
67 print(f"Safe messages: {safe_messages}")
68 return "...", "..."
69
70 stop_token_ids = set()
71 if tokenizer.eos_token_id is not None:
72 stop_token_ids.add(tokenizer.eos_token_id)
73
74 for token in ["<|im_end|>", "<|endoftext|>"]:
75 tid = tokenizer.convert_tokens_to_ids(token)
76 if tid is not None and tid != tokenizer.unk_token_id:
77 stop_token_ids.add(tid)
78
79 generated_ids = []
80 last_tokens_buffer = []
81 in_thinking = False
82
83 with torch.no_grad():
84 past_key_values = None
85
86 for _ in range(max_new_tokens):
87 outputs = model(
88 input_ids=input_ids if past_key_values is None else input_ids[:, -1:],
89 past_key_values=past_key_values,
90 use_cache=True
91 )
92 logits = outputs.logits[:, -1, :]
93 past_key_values = outputs.past_key_values
94
95 if repetition_penalty != 1.0 and len(generated_ids) > 0:
96 for token_id in set(generated_ids):
97 if logits[0, token_id] < 0:
98 logits[0, token_id] *= repetition_penalty
99 else:
100 logits[0, token_id] /= repetition_penalty
101
102 logits = logits / temperature
103
104 if top_p < 1.0:
105 sorted_logits, sorted_indices = torch.sort(logits, descending=True)
106 cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
107 remove = cum_probs > top_p
108 remove[..., 1:] = remove[..., :-1].clone()
109 remove[..., 0] = False
110 indices_to_remove = sorted_indices[remove]
111 logits[0, indices_to_remove] = float('-inf')
112
113 probs = torch.softmax(logits, dim=-1)
114 next_token = torch.multinomial(probs, 1)
115 token_id = next_token.item()
116
117 last_tokens_buffer.append(token_id)
118 if len(last_tokens_buffer) > 6:
119 last_tokens_buffer.pop(0)
120 if len(set(last_tokens_buffer)) == 1:
121 break
122
123 input_ids = next_token
124 generated_ids.append(token_id)
125
126 if token_id in stop_token_ids:
127 break
128
129 word = tokenizer.decode([token_id], skip_special_tokens=False)
130
131 if "<think>" in word:
132 in_thinking = True
133 continue
134 if "</think>" in word:
135 in_thinking = False
136 continue
137
138 if not in_thinking:
139 clean_word = word.replace("<|im_end|>", "").replace("<|endoftext|>", "")
140 if clean_word:
141 print(clean_word, end="", flush=True)
142
143 print()
144
145 full_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
146
147 clean_text = re.sub(r'<think>.*?</think>', '', full_text, flags=re.DOTALL).strip()
148 clean_text = clean_text.replace('</think>', '').replace('<think>', '').strip()
149
150 clean_text = sanitize_message(clean_text)
151
152 return full_text, clean_text
153
154
155def main():
156 tokenizer, model = load_model(model_path)
157
158 system_prompt = "Ты общительный участник чата"
159 history = [{"role": "system", "content": system_prompt}]
160
161 while True:
162 try:
163 user_input = input("\nYou: ").strip()
164
165 if not user_input:
166 continue
167 if user_input.lower() in ["exit", "quit"]:
168 break
169 if user_input.lower() == "clear":
170 history = [{"role": "system", "content": system_prompt}]
171 print("History cleared")
172 continue
173
174 history.append({"role": "user", "content": user_input})
175
176 print("Model: ", end="", flush=True)
177
178 full, answer = generate_stream(tokenizer, model, history)
179
180 answer = sanitize_message(answer)
181
182 history.append({"role": "assistant", "content": answer})
183
184 except KeyboardInterrupt:
185 print("\n[Aborted]")
186 if history and history[-1].get('role') == 'user':
187 history.pop()
188 continue
189 except Exception as e:
190 print(f"\nError: {e}")
191 import traceback
192 traceback.print_exc()
193 if history and history[-1].get('role') == 'user':
194 history.pop()
195
196
197if __name__ == "__main__":
198 main()
