Ai-Exocore / test_transformers.py
ChoruYt's picture
Upload 22 files
dfa523b verified
Raw
History Blame Contribute Delete
1.08 kB
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_path = "/home/runner/workspace/model"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
print("Loading model (Qwen3-0.6B)...")
model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.float32,
device_map="cpu",
)
model.eval()
prompt = "Who is Michael Jakson"
print(f"\nPrompt: {prompt}")
print("Generating...\n")
messages = [{"role": "user", "content": prompt}]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=False,
)
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=100,
temperature=0.7,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
)
response = tokenizer.decode(
outputs[0][inputs["input_ids"].shape[1]:],
skip_special_tokens=True,
)
print("Response:", response)
print("\n--- Model loaded and working smoothly ---")