import torch from torch import nn from torch.nn import functional as F import transformers
base_path = 'sanagnos/galactica-6.7b-finetuned'
model = transformers.OPTForCausalLM.from_pretrained( base_path, load_in_8bit=True, device_map='auto', low_cpu_mem_usage=True, torch_dtype=torch.float16, offload_state_dict=True )
model.gradient_checkpointing_enable() # reduce number of stored activations model.model.decoder.project_in = lambda x: x.requires_grad_(True)
class CastOutputToFloat(nn.Sequential): def forward(self, x): return super().forward(x).to(torch.float32)
model.lm_head = CastOutputToFloat(model.lm_head)
tokenizer = transformers.AutoTokenizer.from_pretrained(base_path)
batch = "What are the symptoms of Alzheimer's disease?"
with torch.cuda.amp.autocast(): out = model.generate( input_ids=batch['input_ids'], max_length=300, do_sample=True, top_k=40, num_beams=1, num_return_sequences=1, eos_token_id=tokenizer.additional_special_tokens_ids[tokenizer.additional_special_tokens.index('')] )
message = tokenizer.decode(out[0, :-1]).replace('', "User:\n").replace('', 'Assistant:\n')