Buckets:
| from typing import Dict, Any, Tuple, Optional | |
| import torch | |
| from peft import get_peft_model, LoraConfig, PeftModel, TaskType | |
| from transformers import BitsAndBytesConfig | |
| class LoRATTTAdapter: | |
| def __init__(self, model_args): | |
| self.model_args = model_args | |
| self.device = model_args.device | |
| self.lora_rank = model_args.lora_rank or 8 | |
| self.lora_alpha = model_args.lora_alpha or 16 | |
| self.max_length = model_args.max_length or 128 | |
| def setup_lora_model(self, base_model) -> Tuple[PeftModel, LoraConfig]: | |
| lora_config = LoraConfig( | |
| r=self.lora_rank, | |
| lora_alpha=self.lora_alpha, | |
| target_modules=["c_attn"], | |
| lora_dropout=0.1, | |
| bias="none", | |
| task_type=TaskType.CAUSAL_LM, | |
| init_lora_weights=True, | |
| ) | |
| lora_model = get_peft_model(model=base_model, peft_config=lora_config) | |
| lora_model.train() | |
| return lora_model, lora_config | |
| def perform_ttt_step( | |
| self, | |
| model: PeftModel, | |
| tokenizer, | |
| prompt: str, | |
| target_output: str, | |
| optimizer: torch.optim.Optimizer, | |
| temperature: float = 0.1, | |
| ) -> float: | |
| inputs = tokenizer( | |
| prompt + " " + target_output, | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=self.max_length, | |
| padding=True, | |
| ) | |
| inputs = {k: v.to(self.device) for k, v in inputs.items()} | |
| optimizer.zero_grad() | |
| outputs = model(**inputs, labels=inputs["input_ids"]) | |
| loss = outputs.loss | |
| loss.backward() | |
| optimizer.step() | |
| return loss.item() | |
Xet Storage Details
- Size:
- 1.67 kB
- Xet hash:
- e798507f1a593ba71d04fafd3a19c764c08c855891c31dedcef8a32b9fd09fdc
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.