import torchfrom transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfigfrom peft import PeftModelMODEL_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# ================= Inference ==============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()"""