lumicero/qwen3-astroph-lora

!pip install -U bitsandbytes transformers peft accelerate

import re, torch
from typing import List, Optional
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import BitsAndBytesConfig
from peft import PeftModel, PeftConfig

_THINK_RE     = re.compile(r"(?is)<think>.*?</think>")
_LEAD_ASSIST  = re.compile(r"(?is)^.*?\bassistant\b[:\-]?\s*")

def clean_text(s: str) -> str:
    s = _THINK_RE.sub("", s)
    s = _LEAD_ASSIST.sub("", s.strip(), count=1)
    return s.strip()

def load_model_and_tokenizer(
    base_model: Optional[str],
    adapter_repo: str,
    max_seq_len: int = 4096,
    load_in_4bit: bool = True,
    compute_dtype: torch.dtype = torch.bfloat16,
):
    peft_cfg = PeftConfig.from_pretrained(adapter_repo)
    suggested_base = getattr(peft_cfg, "base_model_name_or_path", None)

    if base_model is None:
        base_model = suggested_base
        print(f"[info] Using base from adapter config: {base_model}")
    elif suggested_base and (base_model != suggested_base):
        print(f"[warn] Adapter expects base '{suggested_base}', "
              f"but you set '{base_model}'. Make sure they match!")

    quant_cfg = None
    if load_in_4bit:
        quant_cfg = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=compute_dtype,
            bnb_4bit_quant_type="nf4",
        )

    tok = AutoTokenizer.from_pretrained(base_model, use_fast=True, trust_remote_code=True)

    if tok.pad_token_id is None and tok.eos_token_id is not None:
        tok.pad_token = tok.eos_token
        tok.pad_token_id = tok.eos_token_id

    tok.padding_side    = "left"
    tok.truncation_side = "left"

    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        device_map="auto",
        torch_dtype="auto" if not load_in_4bit else None,
        quantization_config=quant_cfg,
        trust_remote_code=True,
    )

    model = PeftModel.from_pretrained(model, adapter_repo)

    if getattr(model, "generation_config", None) is not None and tok.pad_token_id is not None:
        model.generation_config.pad_token_id = tok.pad_token_id

    return model, tok

def build_chat(tok, user_text: str, system_text: Optional[str] = None) -> str:
    messages = []
    if system_text:
        messages.append({"role": "system", "content": system_text})
    messages.append({"role": "user", "content": user_text})
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    return prompt

@torch.inference_mode()
def generate_one(
    model,
    tok,
    user_text: str,
    system_text: str = "Answer concisely, straight to the point, no <think>.",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    top_p: float = 0.9,
):
    prompt = build_chat(tok, user_text, system_text)
    device = next(model.parameters()).device
    inputs = tok(prompt, return_tensors="pt").to(device)
    in_len = inputs["input_ids"].shape[1]

    im_end_id = tok.convert_tokens_to_ids("<|im_end|>")
    eos_ids = [i for i in [tok.eos_token_id, im_end_id] if i is not None]
    eos_ids = eos_ids[0] if len(eos_ids) == 1 else eos_ids

    out = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=(temperature is not None and temperature > 0),
        temperature=temperature,
        top_p=top_p,
        num_beams=1,
        eos_token_id=eos_ids,
        pad_token_id=tok.pad_token_id,
        use_cache=True,
        no_repeat_ngram_size=3,
        repetition_penalty=1.1,
    )[0]

    gen_ids = out[in_len:]
    text = tok.decode(gen_ids, skip_special_tokens=True)
    return clean_text(text)

@torch.inference_mode()
def generate_batch(
    model,
    tok,
    user_texts: List[str],
    system_text: Optional[str] = None,
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    top_p: float = 0.9,
    batch_size: int = 8,
):
    device = next(model.parameters()).device
    im_end_id = tok.convert_tokens_to_ids("<|im_end|>")
    eos_ids = [i for i in [tok.eos_token_id, im_end_id] if i is not None]
    eos_ids = eos_ids[0] if len(eos_ids) == 1 else eos_ids

    answers = []
    for i in range(0, len(user_texts), batch_size):
        chunk = user_texts[i:i + batch_size]
        prompts = [build_chat(tok, u, system_text) for u in chunk]
        toks = tok(prompts, return_tensors="pt", padding=True).to(device)
        in_lens = toks["attention_mask"].sum(dim=1).tolist()

        outs = model.generate(
            **toks,
            max_new_tokens=max_new_tokens,
            do_sample=(temperature is not None and temperature > 0),
            temperature=temperature,
            top_p=top_p,
            num_beams=1,
            eos_token_id=eos_ids,
            pad_token_id=tok.pad_token_id,
            use_cache=True,
            no_repeat_ngram_size=3,
            repetition_penalty=1.1,
        )

        for out, L in zip(outs, in_lens):
            ans = tok.decode(out[L:], skip_special_tokens=True)
            answers.append(clean_text(ans))
    return answers

if __name__ == "__main__":
    adapter = "luminolous/astropher-lora"
    base = "unsloth/Qwen3-1.7B-unsloth-bnb-4bit"

    model, tok = load_model_and_tokenizer(
        base_model=base,
        adapter_repo=adapter,
        max_seq_len=4096,
        load_in_4bit=True,
        compute_dtype=torch.bfloat16,
    )

    q = "What is inside a black hole?" # <- You can change the question here
    print(f"\nModel output: {generate_one(model, tok, q, max_new_tokens=180)}")