初始化项目,由ModelHub XC社区提供模型
Model: g023/Qwen3-1.77B-g023-NF4 Source: Original Platform
This commit is contained in:
65
model.py
Normal file
65
model.py
Normal file
@@ -0,0 +1,65 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user