mirror of
https://github.com/zenlm/zen-eco.git
synced 2026-10-30 14:58:11 +00:00
- Apache 2.0 License - Zen Research DAO under Zoo Labs Inc (501(c)(3)) - San Francisco-based non-profit for open AI research - Copyright 2025 Zen Research Authors
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
#!/usr/bin/env python3.13
|
|
"""Test agent identity"""
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
import torch
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"/Users/z/work/zen/zen-eco/agent/finetuned",
|
|
torch_dtype=torch.float16,
|
|
device_map="auto",
|
|
trust_remote_code=True
|
|
)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
"/Users/z/work/zen/zen-eco/agent/finetuned",
|
|
trust_remote_code=True
|
|
)
|
|
|
|
prompts = ["Who are you?", "What model are you?", "What are your capabilities?"]
|
|
|
|
for prompt in prompts:
|
|
text = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
|
|
inputs = tokenizer(text, return_tensors="pt")
|
|
inputs = {k: v.to(model.device) for k, v in inputs.items()}
|
|
|
|
with torch.no_grad():
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=50,
|
|
temperature=0.1,
|
|
do_sample=False,
|
|
pad_token_id=tokenizer.eos_token_id
|
|
)
|
|
|
|
response = tokenizer.decode(outputs[0], skip_special_tokens=False)
|
|
print(f"Q: {prompt}")
|
|
print(f"A: {response.split('assistant')[-1].strip()}")
|
|
print() |