Boogon/sakura-qwen3-0.6b-lora-demo-v1

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load the Qwen3-0.6B chat model and tokenizer
model_name = "Qwen/Qwen3-0.6B"
adapter_name = "Boogon/sakura-qwen3-0.6b-lora-demo-v1"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)
model = PeftModel.from_pretrained(model, adapter_name)

def chat_with_vtuber(messages, max_length=512):
    """
    messages format: [
        {"role": "system", "content": ""},
        {"role": "user", "content": "Hello!"},
        {"role": "assistant", "content": "Hi there, my name's Sakura!(*´▽`*)"},
        {"role": "user", "content": "What's new today?"}
        # ...
    ]
    """
    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=max_length,
            temperature=0.8,
            top_p=0.9,
            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)
    return response

# Example conversation
messages = [
    {"role": "system", "content": "You are Sakura-chan, a cute VTuber who loves interacting with fans. You speak in a cheerful, cute style with occasional Japanese phrases and emojis."},
    {"role": "user", "content": "Hello Sakura-chan! How are you today?"}
]

response = chat_with_vtuber(messages)
print(f"Sakura: {response}")

# Alternative method using chat template directly
conversation = [
    {"role": "user", "content": "What games do you like to play?"}
]

# Apply chat template
text = tokenizer.apply_chat_template(
    conversation,
    tokenize=False,
    add_generation_prompt=True
)

inputs = tokenizer(text, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=256,
    temperature=0.7,
    do_sample=True
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)