| """Model loading and affect-conditioned generation for CODELOCK.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import List, Optional |
| from contextlib import nullcontext |
|
|
| import torch |
|
|
| DEFAULT_MODEL = "Qwen/Qwen3-8B" |
|
|
| from .persona import CODELOCK_SYSTEM |
|
|
|
|
| def load( |
| model_id: str = DEFAULT_MODEL, |
| four_bit: bool = False, |
| device_map: str = "auto", |
| revision=None, |
| ): |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| tokenizer = AutoTokenizer.from_pretrained( |
| model_id, revision=revision, trust_remote_code=False |
| ) |
| kwargs = { |
| "device_map": device_map, |
| "torch_dtype": torch.bfloat16 if torch.cuda.is_available() else torch.float32, |
| "revision": revision, |
| "trust_remote_code": False, |
| "attn_implementation": "sdpa", |
| } |
| if four_bit: |
| from transformers import BitsAndBytesConfig |
|
|
| kwargs["quantization_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_id, **kwargs) |
| model.eval() |
| return model, tokenizer |
|
|
|
|
| def format_prompt(tokenizer, user, history=None, system=CODELOCK_SYSTEM): |
| messages = ( |
| [{"role": "system", "content": system}] |
| + list(history or []) |
| + [{"role": "user", "content": user}] |
| ) |
| kwargs = {"tokenize": False, "add_generation_prompt": True} |
| try: |
| prompt = tokenizer.apply_chat_template( |
| messages, enable_thinking=False, **kwargs |
| ) |
| except TypeError: |
| prompt = tokenizer.apply_chat_template(messages, **kwargs) |
| return prompt |
|
|
|
|
| @torch.no_grad() |
| def generate( |
| model, |
| tokenizer, |
| user: str, |
| history: Optional[List[dict]] = None, |
| system: str = CODELOCK_SYSTEM, |
| max_new_tokens: int = 200, |
| temperature: float = 0.7, |
| top_p: float = 0.8, |
| top_k: int = 20, |
| repetition_penalty: float = 1.15, |
| seed: Optional[int] = None, |
| do_sample: bool = True, |
| ) -> str: |
| prompt = format_prompt(tokenizer, user, history, system) |
| inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to( |
| model.get_input_embeddings().weight.device |
| ) |
| owner = getattr(model, "_codelock_steerer", None) |
| sampling = ( |
| {"temperature": temperature, "top_p": top_p, "top_k": top_k} |
| if do_sample |
| else {} |
| ) |
| devices = sorted( |
| {p.device.index for p in model.parameters() if p.device.type == "cuda"} |
| ) |
| with torch.random.fork_rng(devices=devices), ( |
| owner.frozen() if owner is not None else nullcontext() |
| ): |
| if seed is not None: |
| torch.manual_seed(seed) |
| out = model.generate( |
| **inputs, |
| max_new_tokens=max_new_tokens, |
| do_sample=do_sample, |
| **sampling, |
| repetition_penalty=repetition_penalty, |
| pad_token_id=( |
| tokenizer.pad_token_id |
| if tokenizer.pad_token_id is not None |
| else tokenizer.eos_token_id |
| ), |
| ) |
| return tokenizer.decode( |
| out[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True |
| ).strip() |
|
|
