76 lines
1.5 KiB
Python
76 lines
1.5 KiB
Python
|
import torch
|
||
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
|
|
|
||
|
|
MODEL_PATH = "./outputs/qwen3_0.6b_fft"
|
||
|
|
|
||
|
|
print("Loading tokenizer...")
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
||
|
|
MODEL_PATH,
|
||
|
|
trust_remote_code=True
|
||
|
|
)
|
||
|
|
|
||
|
|
print("Loading model...")
|
||
|
|
|
||
|
|
dtype = (
|
||
|
|
torch.bfloat16
|
||
|
|
if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
|
||
|
|
else torch.float16
|
||
|
|
)
|
||
|
|
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL_PATH,
|
||
|
|
dtype=dtype,
|
||
|
|
trust_remote_code=True,
|
||
|
|
device_map="auto"
|
||
|
|
)
|
||
|
|
|
||
|
|
model.eval()
|
||
|
|
|
||
|
|
print("=" * 60)
|
||
|
|
print("Qwen3 Full Fine-Tuned Model")
|
||
|
|
print("Type 'exit' to quit")
|
||
|
|
print("=" * 60)
|
||
|
|
|
||
|
|
while True:
|
||
|
|
|
||
|
|
user_input = input("\nUser: ")
|
||
|
|
|
||
|
|
if user_input.lower() in ["exit", "quit"]:
|
||
|
|
break
|
||
|
|
|
||
|
|
messages = [
|
||
|
|
{
|
||
|
|
"role": "user",
|
||
|
|
"content": user_input
|
||
|
|
}
|
||
|
|
]
|
||
|
|
|
||
|
|
text = tokenizer.apply_chat_template(
|
||
|
|
messages,
|
||
|
|
tokenize=False,
|
||
|
|
add_generation_prompt=True
|
||
|
|
)
|
||
|
|
|
||
|
|
inputs = tokenizer(
|
||
|
|
text,
|
||
|
|
return_tensors="pt"
|
||
|
|
).to(model.device)
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
|
||
|
|
outputs = model.generate(
|
||
|
|
**inputs,
|
||
|
|
max_new_tokens=512,
|
||
|
|
temperature=0.7,
|
||
|
|
top_p=0.95,
|
||
|
|
do_sample=True,
|
||
|
|
repetition_penalty=1.1,
|
||
|
|
eos_token_id=tokenizer.eos_token_id
|
||
|
|
)
|
||
|
|
|
||
|
|
response = tokenizer.decode(
|
||
|
|
outputs[0][inputs["input_ids"].shape[1]:],
|
||
|
|
skip_special_tokens=True
|
||
|
|
)
|
||
|
|
|
||
|
|
print("\nAssistant:", response)
|