import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
MODEL_NAME = "ahammad115566/smeft-qwen-7b"
RESPONSE_PREFIX = "\n### Response:\n"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
quantization_config=bnb_config,
device_map={"": 0},
trust_remote_code=True,
)
model.eval()
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
def build_prompt(instruction: str, context: str = "") -> str:
parts = [ f"\n### Instruction:\n{instruction}"]
parts.append(RESPONSE_PREFIX)
return "\n".join(parts)
@torch.inference_mode()
def ask(instruction: str, context: str = "") -> str:
inputs = tokenizer(
build_prompt(instruction, context),
return_tensors="pt",
add_special_tokens=True,
).to("cuda:0")
output_ids = model.generate(
**inputs,
max_new_tokens=512,
do_sample=False,
repetition_penalty=1.1,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.pad_token_id,
)
new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
"""