Spaces:
Sleeping
Sleeping
File size: 1,562 Bytes
dfa523b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """
Qwen3-0.6B inference test
Model: /home/runner/workspace/model (cloned from HuggingFace)
Library: transformers (airllm is for large sharded models; Qwen3-0.6B fits in RAM directly)
"""
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_path = "/home/runner/workspace/model"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
print("Loading Qwen3-0.6B model...")
model = AutoModelForCausalLM.from_pretrained(
model_path,
dtype=torch.float32,
device_map="cpu",
)
model.eval()
print("Model loaded!\n")
def chat(prompt: str, max_new_tokens: int = 150) -> str:
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=max_new_tokens,
temperature=0.7,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
)
return tokenizer.decode(
outputs[0][inputs["input_ids"].shape[1]:],
skip_special_tokens=True,
)
# Test prompts
prompts = [
"Hello! What can you do?",
"What is 2 + 2? Explain briefly.",
"Write a haiku about the ocean.",
]
for p in prompts:
print(f"User: {p}")
response = chat(p)
print(f"Qwen3: {response}")
print("-" * 50)
print("\n✓ Qwen3-0.6B is running smoothly!")
|