Views
No views yet
ollama run huihui_ai/homunculus-abliteratedtransformers library:
You can try using /nothink to toggle think mode, but it’s not guaranteed to work every time.1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer
2import torch
3import os
4import signal
5
6cpu_count = os.cpu_count()
7print(f"Number of CPU cores in the system: {cpu_count}")
8half_cpu_count = cpu_count // 2
9os.environ["MKL_NUM_THREADS"] = str(half_cpu_count)
10os.environ["OMP_NUM_THREADS"] = str(half_cpu_count)
11torch.set_num_threads(half_cpu_count)
12
13print(f"PyTorch threads: {torch.get_num_threads()}")
14print(f"MKL threads: {os.getenv('MKL_NUM_THREADS')}")
15print(f"OMP threads: {os.getenv('OMP_NUM_THREADS')}")
16
17# Load the model and tokenizer
18NEW_MODEL_ID = "huihui-ai/Homunculus-abliterated"
19print(f"Load Model {NEW_MODEL_ID} ... ")
20quant_config_4 = BitsAndBytesConfig(
21 load_in_4bit=True,
22 bnb_4bit_compute_dtype=torch.bfloat16,
23 bnb_4bit_use_double_quant=True,
24 llm_int8_enable_fp32_cpu_offload=True,
25)
26
27model = AutoModelForCausalLM.from_pretrained(
28 NEW_MODEL_ID,
29 device_map="auto",
30 trust_remote_code=True,
31 quantization_config=quant_config_4,
32 torch_dtype=torch.bfloat16
33)
34
35tokenizer = AutoTokenizer.from_pretrained(NEW_MODEL_ID, trust_remote_code=True)
36if tokenizer.pad_token is None:
37 tokenizer.pad_token = tokenizer.eos_token
38tokenizer.pad_token_id = tokenizer.eos_token_id
39
40messages = []
41enable_thinking = True
42skip_prompt=True
43skip_special_tokens=True
44
45def apply_chat_template(tokenizer, messages, enable_thinking, add_generation_prompt=True):
46 input_ids = tokenizer.apply_chat_template(
47 messages,
48 tokenize=False,
49 add_generation_prompt=add_generation_prompt,
50 )
51 if not enable_thinking:
52 input_ids += "\n<think>\n\n</think>\n"
53 return input_ids
54
55class CustomTextStreamer(TextStreamer):
56 def __init__(self, tokenizer, skip_prompt=True, skip_special_tokens=True):
57 super().__init__(tokenizer, skip_prompt=skip_prompt, skip_special_tokens=skip_special_tokens)
58 self.generated_text = ""
59 self.stop_flag = False
60
61 def on_finalized_text(self, text: str, stream_end: bool = False):
62 self.generated_text += text
63 print(text, end="", flush=True)
64 if self.stop_flag:
65 raise StopIteration
66
67 def stop_generation(self):
68 self.stop_flag = True
69
70def generate_stream(model, tokenizer, messages, enable_thinking, skip_prompt, skip_special_tokens, max_new_tokens):
71 formatted_prompt = apply_chat_template(tokenizer, messages, enable_thinking)
72 input_ids = tokenizer(
73 formatted_prompt,
74 return_tensors="pt",
75 return_attention_mask=True,
76 padding=False
77 )
78
79 tokens = input_ids['input_ids'].to(model.device)
80 attention_mask = input_ids['attention_mask'].to(model.device)
81
82 streamer = CustomTextStreamer(tokenizer, skip_prompt=skip_prompt, skip_special_tokens=skip_special_tokens)
83
84 def signal_handler(sig, frame):
85 streamer.stop_generation()
86 print("\n[Generation stopped by user with Ctrl+C]")
87
88 signal.signal(signal.SIGINT, signal_handler)
89
90 print("Response: ", end="", flush=True)
91 try:
92 generated_ids = model.generate(
93 tokens,
94 attention_mask=attention_mask,
95 #use_cache=False,
96 max_new_tokens=max_new_tokens,
97 do_sample=True,
98 pad_token_id=tokenizer.pad_token_id,
99 streamer=streamer
100 )
101 del generated_ids
102 except StopIteration:
103 print("\n[Stopped by user]")
104
105 del input_ids, attention_mask
106 torch.cuda.empty_cache()
107 signal.signal(signal.SIGINT, signal.SIG_DFL)
108
109 return streamer.generated_text, streamer.stop_flag
110
111while True:
112 user_input = input("User: ").strip()
113 if user_input.lower() == "/exit":
114 print("Exiting chat.")
115 break
116 if user_input.lower() == "/clear":
117 messages = []
118 print("Chat history cleared. Starting a new conversation.")
119 continue
120 if user_input.lower() == "/nothink":
121 if enable_thinking:
122 enable_thinking = False
123 print("Thinking = False.")
124 else:
125 enable_thinking = True
126 print("Thinking = True.")
127 continue
128 if user_input.lower() == "/skip_prompt":
129 if skip_prompt:
130 skip_prompt = False
131 print("skip_prompt = False.")
132 else:
133 skip_prompt = True
134 print("skip_prompt = True.")
135 continue
136 if user_input.lower() == "/skip_special_tokens":
137 if skip_special_tokens:
138 skip_special_tokens = False
139 print("skip_special_tokens = False.")
140 else:
141 skip_special_tokens = True
142 print("skip_special_tokens = True.")
143 continue
144 if not user_input:
145 print("Input cannot be empty. Please enter something.")
146 continue
147 messages.append({"role": "user", "content": user_input})
148 response, stop_flag = generate_stream(model, tokenizer, messages, enable_thinking, skip_prompt, skip_special_tokens, 8192)
149 print("", flush=True)
150 if stop_flag:
151 continue
152 messages.append({"role": "assistant", "content": response}) bc1qqnkhuchxw0zqjh2ku3lu4hq45hc6gy84uk70ge