65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
# Tweakable parameters
|
|
# MODEL_PATH = "./Qwen3-NF4" # local version
|
|
MODEL_PATH = "g023/Qwen3-1.77B-g023-NF4"
|
|
MAX_NEW_TOKENS = 2000
|
|
TEMPERATURE = 0.7
|
|
DO_SAMPLE = True
|
|
TOP_P = 0.9
|
|
TOP_K = 50
|
|
REPETITION_PENALTY = 1.1
|
|
STREAMING = True # Set to True for streaming inference
|
|
INPUT_MESSAGE = "You are completing the next step in a task to create an arcade game in javascript. Your available tools are rationalize, red_green_tdd, and create_plan. Synthesize their output when reasoning. "
|
|
|
|
import sys
|
|
import os
|
|
sys.path.append(os.path.dirname(__file__))
|
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
|
|
|
|
def load_model():
|
|
print("Loading model...")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_PATH,
|
|
device_map="auto",
|
|
)
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
|
|
print("Model loaded.")
|
|
return model, tokenizer
|
|
|
|
def inference_non_streaming(model, tokenizer, messages):
|
|
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
|
|
inputs = tokenizer(text, return_tensors="pt").to(model.device)
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=MAX_NEW_TOKENS,
|
|
temperature=TEMPERATURE,
|
|
do_sample=DO_SAMPLE,
|
|
top_p=TOP_P,
|
|
top_k=TOP_K,
|
|
repetition_penalty=REPETITION_PENALTY,
|
|
)
|
|
response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
|
print("Response:", response)
|
|
|
|
def inference_streaming(model, tokenizer, messages):
|
|
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
|
|
inputs = tokenizer(text, return_tensors="pt").to(model.device)
|
|
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=MAX_NEW_TOKENS,
|
|
temperature=TEMPERATURE,
|
|
do_sample=DO_SAMPLE,
|
|
top_p=TOP_P,
|
|
top_k=TOP_K,
|
|
repetition_penalty=REPETITION_PENALTY,
|
|
streamer=streamer,
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
model, tokenizer = load_model()
|
|
messages = [{"role": "user", "content": INPUT_MESSAGE}]
|
|
if STREAMING:
|
|
inference_streaming(model, tokenizer, messages)
|
|
else:
|
|
inference_non_streaming(model, tokenizer, messages) |