File size: 2,444 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | import base64
import time
from typing import Any, Callable, Optional
import os
from ..types import MessageList, SamplerBase
import torch
import gc
from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
class LLaMASampler(SamplerBase):
"""
Sample from LLaDA
"""
def __init__(
self,
model_name: str = "meta-llama/Llama-3.1-8B-Instruct",
generate_fn: Optional[Callable] = None,
generation_kwargs: Optional[dict] = None,
kv_cache_masked: Optional[bool] = False,
kv_cache_decoded: Optional[bool] = False,
):
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
).eval().requires_grad_(False)
self.tokenizer = AutoTokenizer.from_pretrained(
model_name, trust_remote_code=True
)
self.generation_kwargs = generation_kwargs
def init_model(self):
self.model = self.model.cuda()
def _handle_text(self, text: str) -> dict[str, Any]:
return {"type": "input_text", "text": text}
def _pack_message(self, role: str, content: Any) -> dict[str, Any]:
return {"role": role, "content": content}
def _free_memory(self):
del self.model
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
import time
time.sleep(10)
def __call__(self, seq_idx: int, message_list: MessageList) -> str:
with torch.inference_mode():
full_message = [
{'role': 'system', 'content': 'You are a helpful AI assistant.'}
] + message_list
prompt = self.tokenizer.apply_chat_template(full_message, add_generation_prompt=True, tokenize=False)
#print(prompt)
model_inputs = self.tokenizer([prompt], return_tensors="pt").to(self.model.device)
generated_ids = self.model.generate(
**model_inputs,
**self.generation_kwargs
)
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
length = len(generated_ids[0])
res = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
#print(response)
return seq_idx, res
|