from unsloth import FastLanguageModel
from peft import PeftModel
from transformers import AutoTokenizer, TextStreamer
BASE_ID = "unsloth/Qwen2.5-1.5B-Instruct"
LORA_REPO = "qailunu/kw5-lite-swahili-lora"
print("Loading KW5-Lite...")
tokenizer = AutoTokenizer.from_pretrained(LORA_REPO)
model, _ = FastLanguageModel.from_pretrained(
model_name=BASE_ID,
max_seq_length=2048,
load_in_4bit=False,
)
model = PeftModel.from_pretrained(model, LORA_REPO)
FastLanguageModel.for_inference(model)
model.generation_config.max_length = None
print("✅ Loaded.\n")
def chat(prompt, history=None, max_new_tokens=300, temperature=0.7, top_p=0.9, stream=True):
if history is None:
history = []
history.append({"role": "user", "content": prompt})
inputs = tokenizer.apply_chat_template(
history,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
).to(model.device)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) if stream else None
output_ids = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
do_sample=True,
streamer=streamer,
pad_token_id=tokenizer.eos_token_id,
)
reply = tokenizer.decode(
output_ids[0][inputs["input_ids"].shape[-1]:],
skip_special_tokens=True,
).strip()
history.append({"role": "assistant", "content": reply})
return reply, history
print("💬 Chat na KW5-Lite (andika 'quit' kutoka)\n")
chat_history = []
while True:
user_input = input("👤 ")
if user_input.strip().lower() in {"quit", "exit", "toka"}:
break
print("🤖 ", end="")
_, chat_history = chat(user_input, chat_history)
print("\n" + "-" * 40)