inference.py

import torch
from sentence_transformers import SentenceTransformer
from hypernet import HyperNet

def inject_hyperlora(text_prompt, hypernet_path, target_shapes, r=8, device="cuda"):
    """
    Phase 5: Inference & Assembly
    Takes a new unseen text prompt, passes through HyperNet, maps to A and B representation,
    and returns the assembled standard LoRA matrices ready for injection.
    """
    print(f"Embedding text prompt: '{text_prompt}'")
    encoder = SentenceTransformer("sentence-transformers/all-mpnet-base-v2").to(device)
    encoder.eval()
    
    with torch.no_grad():
        z = encoder.encode([text_prompt], convert_to_tensor=True).to(device)
        
    hypernet = HyperNet(target_shapes, r=r).to(device)
    hypernet.load_state_dict(torch.load(hypernet_path, map_location=device))
    hypernet.eval()
    
    print("Generating SVD components from Hypernetwork...")
    with torch.no_grad():
        U_gen_dict, log_Sigma_gen_dict, V_gen_dict = hypernet(z)
        
    assembled_matrices = {}
    
    for name in U_gen_dict.keys():
        U_gen = U_gen_dict[name].squeeze(0)          # (d_out, r)
        log_Sigma_gen = log_Sigma_gen_dict[name].squeeze(0)  # (r)
        V_gen_T = V_gen_dict[name].squeeze(0)        # (r, d_in) -> represents V^T
        
        # Reverse the log to strictly positive
        Sigma_gen = torch.exp(log_Sigma_gen)
        sqrt_diag_Sigma = torch.diag(torch.sqrt(Sigma_gen))
        
        # Assemble standard LoRA matrices
        # A = U_gen x sqrt(diag(Sigma_gen))
        A = torch.matmul(U_gen, sqrt_diag_Sigma)     # (d_out, r)
        
        # B = sqrt(diag(Sigma_gen)) x V_gen^T
        B = torch.matmul(sqrt_diag_Sigma, V_gen_T)   # (r, d_in)
        
        assembled_matrices[name] = {
            "A": A,
            "B": B
        }
    
    print(f"Successfully generated standard LoRA matrices for {len(assembled_matrices)} modules.")
    return assembled_matrices

if __name__ == "__main__":
    # Example usage / smoke test
    dummy_shapes = {
        "q_proj": (4096, 4096),
        "v_proj": (1024, 4096)
    }
    # This assumes we have a trained hypernet named hypernet_final.pt
    # matrices = inject_hyperlora("Perform mathematics instruction tuning", "hypernet_final.pt", dummy_shapes)