import re
import json
import torch
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
TextStreamer,
StoppingCriteria,
StoppingCriteriaList,
)
MODEL_NAME = "KordAI/KeawGPT"
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
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)
history = [{"role": "user", "content": "What's the difference between a P and an NP problem?"}]
prompt = build_prompt(history)
output = generate(prompt)
print(output)