Spaces:
Running
Running
| """Model core for Code2LoRA-GRU inference. | |
| Faithful port of ``hypernetwork/code2lora_core.py`` from the RepoPeftData | |
| project, trimmed to the pieces needed for inference (the parquet loaders that | |
| the trainer uses are intentionally omitted so this Space has no pyarrow | |
| dependency). | |
| The two networks that get loaded from a ``gru_head.*.pt`` checkpoint are: | |
| * :class:`CommitGRU` -- a 1-layer GRU that ingests per-commit diff | |
| embeddings; ``h_0`` comes from a learnable projection of the initial | |
| repo-state embedding. | |
| * :class:`Code2LoRAHead` -- maps the GRU hidden state to one LoRA (A, B) | |
| pair per module *type* (q_proj, v_proj, ...), shared across all | |
| transformer layers. | |
| The generated weights are pushed into :class:`LoRA` wrappers that replace the | |
| target ``nn.Linear`` modules of the frozen base LLM. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import re | |
| from dataclasses import dataclass | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| # --------------------------------------------------------------------------- | |
| # LoRA module + injection | |
| # --------------------------------------------------------------------------- | |
| class LoRA(nn.Module): | |
| """Wraps an ``nn.Linear`` with an additive low-rank update. | |
| Forward: ``y = base(x) + scaling * (x @ A^T) @ B^T`` where the per-call | |
| A: ``[rank, in_features]`` and B: ``[out_features, rank]`` come from the | |
| hypernet via :meth:`set_lora_weights`. | |
| """ | |
| def __init__(self, base: nn.Linear, in_features: int, out_features: int, | |
| rank: int, alpha: float): | |
| super().__init__() | |
| self.base = base | |
| for p in self.base.parameters(): | |
| p.requires_grad = False | |
| self.in_features = in_features | |
| self.out_features = out_features | |
| self.rank = rank | |
| self.scaling = float(alpha) / float(max(1, rank)) | |
| self.A: Optional[torch.Tensor] = None # [rank, in_features] | |
| self.B: Optional[torch.Tensor] = None # [out_features, rank] | |
| def set_lora_weights(self, A: torch.Tensor, B: torch.Tensor) -> None: | |
| self.A = A | |
| self.B = B | |
| def clear_lora_weights(self) -> None: | |
| self.A = None | |
| self.B = None | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| y = self.base(x) | |
| if self.A is None or self.B is None: | |
| return y | |
| x_f32 = x.detach().to(torch.float32) | |
| A = self.A.to(torch.float32) | |
| B = self.B.to(torch.float32) | |
| delta = F.linear(F.linear(x_f32, A), B) * self.scaling | |
| return y + delta.to(dtype=y.dtype) | |
| class ModuleSpec: | |
| full_name: str # e.g. 'model.layers.5.self_attn.q_proj' | |
| layer_idx: int | |
| type: str # e.g. 'q_proj' | |
| in_features: int | |
| out_features: int | |
| def get_module_specs(model: nn.Module, target_module_types: List[str] | |
| ) -> List[ModuleSpec]: | |
| """Discover every nn.Linear whose name contains a target type.""" | |
| specs: List[ModuleSpec] = [] | |
| pat = re.compile(r"\bmodel\.layers\.(\d+)\.") | |
| for name, m in model.named_modules(): | |
| match_type = next((t for t in target_module_types if t in name), None) | |
| if match_type is None: | |
| continue | |
| if not isinstance(m, nn.Linear): | |
| continue | |
| m_layer = pat.search(name) | |
| layer_idx = int(m_layer.group(1)) if m_layer else -1 | |
| specs.append(ModuleSpec( | |
| full_name=name, | |
| layer_idx=layer_idx, | |
| type=match_type, | |
| in_features=int(m.in_features), | |
| out_features=int(m.out_features), | |
| )) | |
| specs.sort(key=lambda s: (s.layer_idx, s.full_name)) | |
| return specs | |
| def replace_with_lora(model: nn.Module, specs: List[ModuleSpec], | |
| rank: int, alpha: float) -> None: | |
| """Replace each target ``nn.Linear`` with a :class:`LoRA` wrapper. Idempotent.""" | |
| named = dict(model.named_modules()) | |
| device = next(model.parameters()).device | |
| dtype = next(model.parameters()).dtype | |
| for sp in specs: | |
| parent_name, attr = sp.full_name.rsplit(".", 1) | |
| orig = getattr(named[parent_name], attr) | |
| if isinstance(orig, LoRA): | |
| continue | |
| assert isinstance(orig, nn.Linear), \ | |
| f"{sp.full_name} is not nn.Linear (got {type(orig)})" | |
| wrapped = LoRA(orig, sp.in_features, sp.out_features, | |
| rank, alpha).to(device=device, dtype=dtype) | |
| setattr(named[parent_name], attr, wrapped) | |
| def inject_lora_weights(model: nn.Module, specs: List[ModuleSpec], | |
| head_out: Dict[str, Dict[str, torch.Tensor]], | |
| batch_index: int = 0) -> None: | |
| """Push ``head_out['A'][type]`` / ``head_out['B'][type]`` into each LoRA.""" | |
| A_by_type = head_out["A"] | |
| B_by_type = head_out["B"] | |
| named = dict(model.named_modules()) | |
| for sp in specs: | |
| named[sp.full_name].set_lora_weights( | |
| A_by_type[sp.type][batch_index], | |
| B_by_type[sp.type][batch_index], | |
| ) | |
| def clear_all_lora_weights(model: nn.Module, specs: List[ModuleSpec]) -> None: | |
| """Reset every wrapper LoRA to identity (base model only).""" | |
| named = dict(model.named_modules()) | |
| for sp in specs: | |
| mod = named.get(sp.full_name) | |
| if isinstance(mod, LoRA): | |
| mod.clear_lora_weights() | |
| def discover_module_types_and_dims(specs: List[ModuleSpec] | |
| ) -> Dict[str, Tuple[int, int]]: | |
| type_dims: Dict[str, Tuple[int, int]] = {} | |
| for sp in specs: | |
| if sp.type in type_dims: | |
| assert type_dims[sp.type] == (sp.in_features, sp.out_features), \ | |
| f"type {sp.type} appears with inconsistent dims" | |
| continue | |
| type_dims[sp.type] = (sp.in_features, sp.out_features) | |
| return type_dims | |
| # --------------------------------------------------------------------------- | |
| # Shared LoRA generation head | |
| # --------------------------------------------------------------------------- | |
| class Code2LoRAHead(nn.Module): | |
| """Maps a context vector to one LoRA (A, B) pair per module type.""" | |
| def __init__( | |
| self, | |
| input_dim: int, | |
| type_dims: Dict[str, Tuple[int, int]], | |
| hidden_dim: int = 1024, | |
| rank: int = 16, | |
| init_log_scale: float = -3.5, | |
| ): | |
| super().__init__() | |
| self.input_dim = input_dim | |
| self.hidden_dim = hidden_dim | |
| self.rank = rank | |
| self.type_dims = dict(type_dims) | |
| self.types = sorted(type_dims.keys()) | |
| self.trunk = nn.Sequential( | |
| nn.Linear(input_dim, hidden_dim), | |
| nn.GELU(), | |
| nn.Linear(hidden_dim, hidden_dim), | |
| nn.GELU(), | |
| ) | |
| self.heads_A = nn.ModuleDict({ | |
| t: nn.Linear(hidden_dim, rank * type_dims[t][0]) for t in self.types | |
| }) | |
| self.heads_B = nn.ModuleDict({ | |
| t: nn.Linear(hidden_dim, type_dims[t][1] * rank) for t in self.types | |
| }) | |
| self.log_scale_A = nn.ParameterDict({ | |
| t: nn.Parameter(torch.tensor(init_log_scale)) for t in self.types | |
| }) | |
| self.log_scale_B = nn.ParameterDict({ | |
| t: nn.Parameter(torch.tensor(init_log_scale)) for t in self.types | |
| }) | |
| def forward(self, ctx: torch.Tensor) -> Dict[str, Dict[str, torch.Tensor]]: | |
| if ctx.dim() == 3: | |
| ctx = torch.max(ctx, dim=1).values | |
| h = self.trunk(ctx.float()) | |
| h = F.normalize(h, p=2, dim=-1) * math.sqrt(self.hidden_dim) | |
| A_out: Dict[str, torch.Tensor] = {} | |
| B_out: Dict[str, torch.Tensor] = {} | |
| for t in self.types: | |
| in_f, out_f = self.type_dims[t] | |
| A_raw = self.heads_A[t](h).view(-1, self.rank, in_f) | |
| B_raw = self.heads_B[t](h).view(-1, out_f, self.rank) | |
| scale_A = torch.exp(self.log_scale_A[t]).clamp(1e-5, 0.3) | |
| scale_B = torch.exp(self.log_scale_B[t]).clamp(1e-5, 0.3) | |
| A_out[t] = torch.tanh(A_raw) * scale_A | |
| B_out[t] = torch.tanh(B_raw) * scale_B | |
| return {"A": A_out, "B": B_out} | |
| def config_dict(self) -> Dict[str, Any]: | |
| return { | |
| "input_dim": self.input_dim, | |
| "hidden_dim": self.hidden_dim, | |
| "rank": self.rank, | |
| "types": self.types, | |
| "type_dims": {t: list(v) for t, v in self.type_dims.items()}, | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Commit-streaming GRU | |
| # --------------------------------------------------------------------------- | |
| class CommitGRU(nn.Module): | |
| """One-layer GRU that ingests per-commit diff embeddings. | |
| ``h_0`` is a learnable projection of the initial repo-state embedding; | |
| ``h_T`` is fed into :class:`Code2LoRAHead`. | |
| """ | |
| def __init__( | |
| self, | |
| diff_input_dim: int = 2048, | |
| repo_state_dim: int = 2048, | |
| hidden_dim: int = 2048, | |
| ): | |
| super().__init__() | |
| self.diff_input_dim = diff_input_dim | |
| self.repo_state_dim = repo_state_dim | |
| self.hidden_dim = hidden_dim | |
| self.diff_proj = nn.Sequential( | |
| nn.Linear(diff_input_dim, hidden_dim), | |
| nn.LayerNorm(hidden_dim), | |
| ) | |
| self.repo_init_proj = nn.Sequential( | |
| nn.Linear(repo_state_dim, hidden_dim), | |
| nn.GELU(), | |
| nn.LayerNorm(hidden_dim), | |
| ) | |
| self.gru = nn.GRU( | |
| input_size=hidden_dim, | |
| hidden_size=hidden_dim, | |
| num_layers=1, | |
| batch_first=True, | |
| ) | |
| self.output_norm = nn.LayerNorm(hidden_dim) | |
| def init_hidden(self, repo_state_emb_0: torch.Tensor) -> torch.Tensor: | |
| h0 = self.repo_init_proj(repo_state_emb_0.float()) | |
| return h0.unsqueeze(0) | |
| def step(self, diff_emb: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor: | |
| x = self.diff_proj(diff_emb.float()).unsqueeze(1) # [B, 1, H] | |
| _, h_new = self.gru(x, h_prev) | |
| return h_new | |
| def context(self, h: torch.Tensor) -> torch.Tensor: | |
| """Match the eval path: ctx = output_norm(h[-1]).""" | |
| return self.output_norm(h[-1]) | |
| __all__ = [ | |
| "LoRA", | |
| "ModuleSpec", | |
| "get_module_specs", | |
| "replace_with_lora", | |
| "inject_lora_weights", | |
| "clear_all_lora_weights", | |
| "discover_module_types_and_dims", | |
| "Code2LoRAHead", | |
| "CommitGRU", | |
| ] | |