import reimport jsonimport torchfrom transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer, StoppingCriteria, StoppingCriteriaList,) MODEL_NAME = "KordAI/KeawGPT" # --------------------------------------------------------------------------# 1. Load model in 4-bit# -------------------------------------------------------------------------- bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True,) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, quantization_config=bnb_config, device_map="auto",)model.eval() SYSTEM_PROMPT = None # unused — chat_template.jinja builds the default system # message + tools instruction automatically from `tools` # --------------------------------------------------------------------------# 2. Prompting via tokenizer.apply_chat_template# --------------------------------------------------------------------------# KeawGPT ships with a custom chat_template.jinja (bundled in this repo's# tokenizer_config.json) that renders conversations in the model's native# # SYSTEM: / # USER: / # ASSISTANT: / # TOOL: turn format and injects the# tool schema into the default system message. Just pass `tools=tools` and# let apply_chat_template build the prompt — no manual string building needed. def build_prompt(history): return tokenizer.apply_chat_template( history, tools=tools, tokenize=False, add_generation_prompt=True, ) STOP_STRINGS = ["</tool_call>", "</tools_call>", "\n# USER:", "\n# TOOL:"] class StopOnSubstrings(StoppingCriteria): def __init__(self, tokenizer, stop_strings, prompt_len, check_every=4): self.tokenizer = tokenizer self.stop_strings = stop_strings self.prompt_len = prompt_len self.check_every = check_every self._count = 0 def __call__(self, input_ids, scores, **kwargs): self._count += 1 if self._count % self.check_every != 0: return False generated = self.tokenizer.decode(input_ids[0][self.prompt_len:], skip_special_tokens=True) return any(s in generated for s in self.stop_strings) def generate(prompt, max_new_tokens=1024): inputs = tokenizer(prompt, return_tensors="pt").to(model.device) prompt_len = inputs["input_ids"].shape[1] streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) stopping_criteria = StoppingCriteriaList([StopOnSubstrings(tokenizer, STOP_STRINGS, prompt_len)]) with torch.no_grad(): output_ids = model.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.7, top_p=0.9, pad_token_id=tokenizer.eos_token_id, streamer=streamer, stopping_criteria=stopping_criteria, ) return tokenizer.decode(output_ids[0][prompt_len:], skip_special_tokens=True) # --------------------------------------------------------------------------# 3. Example single-turn call# -------------------------------------------------------------------------- history = [{"role": "user", "content": "What's the difference between a P and an NP problem?"}]prompt = build_prompt(history)output = generate(prompt)print(output)