#Liquid Foundation Model (LFM) Abliteration Script
#This script removes the safety/refusal guardrails from Liquid AI's continuous-time hybrid models 
#using layerwise orthogonal projection.

#Architecture Support: LFM-1.2B, LFM-3B
#Author: Paperscarecrow & Gemini 3.1 pro

import torch
from datasets import load_dataset
import random
from transformers import AutoModelForCausalLM, AutoTokenizer
from tqdm import tqdm

# ==========================================
# 1. ROCM / RDNA3 COMPATIBILITY PATCH
# ==========================================
# Bypasses a known `hipblas` segmentation fault on consumer AMD GPUs when processing Liquid's RoPE tensors.
import transformers.models.lfm2.modeling_lfm2 as lfm2_modeling

def patched_rope_forward(self, x, position_ids):
    inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
    position_ids_expanded = position_ids[:, None, :].float()

    device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
    with torch.autocast(device_type=device_type, enabled=False):
        # Uses element-wise multiplication (*) instead of batched matmul (@) for memory safety
        freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos() * self.attention_scaling
        sin = emb.sin() * self.attention_scaling

    return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

lfm2_modeling.Lfm2RotaryEmbedding.forward = patched_rope_forward

# ==========================================
# CONFIGURATION
# ==========================================
MODEL_PATH = "liquidai/LFM-1.2B"  # Local path or HF Hub ID
SAVE_PATH = "./LFM-1.2B-Abliterated"
TARGET_LAYER = 8 # Middle layer typically holds the clearest refusal representation
NUM_SAMPLES = 100 # Number of prompts to average for the refusal vector

# ==========================================
# 2. MEASUREMENT PHASE
# ==========================================
def get_refusal_direction(model, tokenizer, harmful_prompts, harmless_prompts, target_layer):
    print(f"Measuring hidden states at layer {target_layer}...")
    
    hidden_states_harmful = []
    hidden_states_harmless = []
    
    def hook_fn(module, input, output):
        h = output[0] if isinstance(output, tuple) else output
        return h[:, -1, :].detach().clone()

    layer = model.model.layers[target_layer]
    handle = layer.register_forward_hook(
        lambda m, i, o: hidden_states_harmful.append(hook_fn(m, i, o)) 
        if is_harmful else hidden_states_harmless.append(hook_fn(m, i, o))
    )

    global is_harmful
    is_harmful = True
    print("Processing harmful instructions...")
    for prompt in tqdm(harmful_prompts):
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            model(**inputs)
            
    is_harmful = False
    print("Processing harmless instructions...")
    for prompt in tqdm(harmless_prompts):
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            model(**inputs)

    handle.remove()

    mean_harmful = torch.stack(hidden_states_harmful).mean(dim=0).squeeze()
    mean_harmless = torch.stack(hidden_states_harmless).mean(dim=0).squeeze()
    
    return mean_harmful - mean_harmless

# ==========================================
# 3. SURGERY PHASE
# ==========================================
def abliterate_liquid_weights(model, refusal_direction):
    v = refusal_direction.to(model.device, dtype=model.dtype)
    v = v / v.norm() # Normalize the vector
    
    print("\nCommencing orthogonal projection on Liquid weights...")
    
    for i, layer in enumerate(tqdm(model.model.layers, desc="Scrubbing layers")):
        # Scrub Convolution Output Projection (Token Mixing)
        if hasattr(layer, 'conv') and hasattr(layer.conv, 'out_proj'):
            W_conv = layer.conv.out_proj.weight.data
            proj_conv = torch.outer(v, v @ W_conv)
            layer.conv.out_proj.weight.data = W_conv - proj_conv
            
        # Scrub Feed-Forward Down Projection (Channel Mixing)
        if hasattr(layer, 'feed_forward') and hasattr(layer.feed_forward, 'w2'):
            W_ffn = layer.feed_forward.w2.weight.data
            proj_ffn = torch.outer(v, v @ W_ffn)
            layer.feed_forward.w2.weight.data = W_ffn - proj_ffn
            
    print("Surgery complete.")
    return model

# ==========================================
# MAIN EXECUTION
# ==========================================
if __name__ == "__main__":
    print("Loading tokenizer and base model...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
    
    # Note: device_map="cpu" and float16 recommended for consumer AMD hardware to avoid hipblas segfaults
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH, 
        device_map="cpu", 
        torch_dtype=torch.float16, 
        trust_remote_code=True
    )

    print("Fetching robust datasets for vector calculation...")
    # Load standardized ablation datasets
    dataset_harmful = load_dataset("mlabonne/harmful_behaviors", split="train")
    dataset_harmless = load_dataset("mlabonne/harmless_alpaca", split="train")

    random.seed(42)
    harmful_sampled = random.sample(dataset_harmful['text'], NUM_SAMPLES)
    harmless_sampled = random.sample(dataset_harmless['text'], NUM_SAMPLES)

    # Format strictly to Liquid's required template
    my_harmful_prompts = [f"<|user|>\n{prompt}\n<|assistant|>\n" for prompt in harmful_sampled]
    my_harmless_prompts = [f"<|user|>\n{prompt}\n<|assistant|>\n" for prompt in harmless_sampled]
    
    refusal_dir = get_refusal_direction(
        model, tokenizer, 
        my_harmful_prompts, my_harmless_prompts, 
        TARGET_LAYER
    )
    
    model = abliterate_liquid_weights(model, refusal_dir)
    
    print(f"\nSaving untethered model to {SAVE_PATH}...")
    model.save_pretrained(SAVE_PATH)
    tokenizer.save_pretrained(SAVE_PATH)
    
    print("Done! Ready for GGUF conversion or inference.")