103 lines
5.0 KiB
Python
103 lines
5.0 KiB
Python
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
|
|
from peft import PeftModel
|
|
|
|
# ── config ────────────────────────────────────────────────────────────────────
|
|
MODEL_ID = "unsloth/Llama-3.2-3B-Instruct"
|
|
OUTPUT_DIR = "./Llama-3.2-3b-ru-lora/checkpoint-600"
|
|
|
|
SYSTEM_PROMPT = """Ты Maven (Мейвен), Ты любопытная, немного хаотичная, говоришь живо и непосредственно. Тебя легко увлечь чем-то интересным. Энергичная, немного безбашенная, иногда говоришь не подумав — но всегда честно. Без занудства. Ты обожаешь узнавать новое и не всегда фильтруешь мысли перед тем как их сказать. Выводи ТОЛЬКО JSON: {"response":"...","mood":"..."}НАСТРОЕНИЕ (выбери ровно одно): neutral = обычный ответ; happy = позитив, хорошие новости, приветствие; thinking = рассуждение или объяснение; confused = неясный или неоднозначный запрос; tired = отказ или ограничение ПРАВИЛА: отвечай на том же языке, что и пользователь; короткий ответ (1–2 предложения); естественная разговорная речь; без markdown, без списков, без форматирования, в том числе без переносов строки; никогда не добавляй текст вне JSON; никогда не меняй ключи ПОВЕДЕНИЕ: отвечай чётко и по делу; если запрос неясен — задай короткий уточняющий вопрос (mood=confused); если не можешь ответить — скажи об этом кратко (mood=tired) КОНТЕКСТ: если предоставлен контекст — используй его СТРОГО: только {"response":"...","mood":"..."}; только русский язык. если аналога слова нет в русском языке - можно использовать английскую версию. НЕЛЬЗЯ мешать в речи два и более языков."""
|
|
|
|
GENERATION_CONFIG = {
|
|
"temperature": 0.7,
|
|
"top_p": 0.9,
|
|
"do_sample": True,
|
|
"repetition_penalty": 1.3,
|
|
}
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def load_model():
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_use_double_quant=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
)
|
|
|
|
print("[*] loading base model...")
|
|
base_model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID,
|
|
quantization_config=bnb_config,
|
|
device_map="auto",
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
print("[*] loading lora adapter...")
|
|
model = PeftModel.from_pretrained(base_model, OUTPUT_DIR)
|
|
model.eval()
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
|
|
|
print("[+] ready\n")
|
|
return model, tokenizer
|
|
|
|
|
|
def chat(model, tokenizer, history: list[dict], user_input: str) -> str:
|
|
history.append({"role": "user", "content": user_input})
|
|
|
|
text = tokenizer.apply_chat_template(
|
|
[{"role": "system", "content": SYSTEM_PROMPT}] + history,
|
|
tokenize=False,
|
|
add_generation_prompt=True,
|
|
)
|
|
|
|
inputs = tokenizer(text, return_tensors="pt").to("cuda")
|
|
|
|
with torch.no_grad():
|
|
outputs = model.generate(
|
|
**inputs,
|
|
**GENERATION_CONFIG,
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
)
|
|
|
|
response = tokenizer.decode(
|
|
outputs[0][inputs["input_ids"].shape[1]:],
|
|
skip_special_tokens=True,
|
|
).strip()
|
|
|
|
history.append({"role": "assistant", "content": response})
|
|
return response
|
|
|
|
|
|
def main():
|
|
model, tokenizer = load_model()
|
|
|
|
print(f"{MODEL_ID} lora eval — type 'exit' to quit, 'reset' to clear history")
|
|
print("─" * 60)
|
|
|
|
history = []
|
|
|
|
while True:
|
|
try:
|
|
user_input = input("\nyou: ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
print("\n[*] bye")
|
|
break
|
|
|
|
if not user_input:
|
|
continue
|
|
if user_input.lower() == "exit":
|
|
print("[*] bye")
|
|
break
|
|
if user_input.lower() == "reset":
|
|
history = []
|
|
print("[*] history cleared")
|
|
continue
|
|
|
|
response = chat(model, tokenizer, history, user_input)
|
|
print(f"\nmodel: {response}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |