pipeline_ambiguous.py

import os
import torch
import numpy as np
import pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import json
import pyarrow as pa
import pyarrow.parquet as pq

model_id = "meta-llama/Llama-3.2-1B" 
layers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]    
outdir = "./results_ambiguous"
corpus = pd.read_csv("official_corpus.csv") # This script can be used only with the official_corpus or test_corpus_with_priming because the others don't fit the data shape required by my insert EOS function
#print(corpus.head())

device = "cpu" # This is where I will load the model
print(f"Device being used: {device}") 
dtype = torch.float32 # Sets the tensor dtype for the model's weights/activations

def org_corpus(): # Function to separate the columns with each language
    neutral = corpus.iloc[:, 0].dropna().astype(str).tolist() # drops empty lines, makes sure those are strings, adds to list
    pt = corpus.iloc[:, 1].dropna().astype(str).tolist() 
    sp = corpus.iloc[:, 2].dropna().astype(str).tolist()
    fr = corpus.iloc[:, 3].dropna().astype(str).tolist()
    print(f"{neutral}\n{pt}\n{sp}\n{fr}")     
    return neutral, pt, sp, fr

neutral, pt, sp, fr = org_corpus() # Calls the function

# Function to introduce a EOS marker so I know what embeddings to not consider later on
def insert_eos(prime, neutral, eos):
    EOS_sentences = [] # Collects the sentences with the new EOS 
    for p, t in zip(prime, neutral): # Iterates over prime and neutral in pairs and stops at the shorter length. p: current primed string. t: current target substring
        p = p.strip() # REmoves any spaces at beginning or end of sentences
        t = t.strip()
        head, sep, tail = p.partition(t) # Splits p around the first occurence of t. head is the part before the match, sep is the matched substring itself (t) and tail is the part after the match

        if sep: # If t is inside of p, insert EOS in the middle
            EOS_sentences.append(f"{head.rstrip()} {eos} {sep}{tail}") # rstrip() oin the head to avoid double spaces
        else: # Keeps alignment 1:1
            EOS_sentences.append(f"{p} {eos} {t}")
    return EOS_sentences 


def set_up():
    configuration = AutoConfig.from_pretrained(model_id)  # Loads the model’s config without weights to look at model structure/architecture
    # Tells me how many transformer layers the model has and its hidden size:
    print(f"Config: layers={configuration.num_hidden_layers}, hidden={configuration.hidden_size}") 

    tok = AutoTokenizer.from_pretrained(model_id)  # Loads the tokenizer associated with the model I am using 
    if tok.pad_token is None: # If the model lacks a PAD token
        tok.pad_token = tok.eos_token # Then use EOS
    print("Tokenizer ok") # Loading is done

# Loads the model with weights
    model = AutoModelForCausalLM.from_pretrained(  
        model_id,
        torch_dtype=dtype, 
        attn_implementation="eager", # Forces eager attention for compatibility
    )
    model.to(device) # Places model on device and tells me how many transformer blocks it has
    print(f"Model loaded on {device}. Transformer blocks:", len(model.model.layers)) 

    model.eval()  # Puts the model in inference mode so forward passes are deterministic. Training-specific layers are gone
    return tok, model

tokenizer, model = set_up() # Calls the function

# Annotates the test and training sets according to the eos token that marks the split
pt = insert_eos(pt, neutral, tokenizer.eos_token)
sp = insert_eos(sp, neutral, tokenizer.eos_token)
fr = insert_eos(fr, neutral, tokenizer.eos_token)
print(pt[0:])
print(sp[0:])
print(fr[0:])

def mean_pool(hidden, attention_mask, start_indices=None): # Converts token embeddings into a single sentence embedding but also ignores priming sentences
    mask = attention_mask.clone().to(hidden.dtype) # Makes a working copy so I don't mutate the caller’s mask and converts to floar to multiply embeddings. Forces bfloat16

    if start_indices is not None: # If there is a priming sentence, I will use this if block
        B, T = mask.shape
        pos = torch.arange(T, device=mask.device).unsqueeze(0).expand(B, T)
        starts = start_indices.to(mask.device).unsqueeze(1) # One start index per sequence          
        include = (pos >= starts).to(mask.dtype) # This includes 1.0 for positions equal or bigger tha start, else is 0.0 It excludes tokens BEFORE start_indices (the priming sentences). Forces bfloat16
        mask = mask * include # Combines padding mask (attention_mask) with my EOS. Result: 1.0 for tokens to include, 0.0 for tokens to not include.
 
    mask3 = mask.unsqueeze(-1) # Expands the mask to multiply it with the hidden states
    summed = (hidden * mask3).sum(1) # Zeros out padded positions, real tokens keep their value
    counts = mask.sum(1).clamp(min=1.0).unsqueeze(-1) # Counts real tokens in the sentence
    average = summed / counts # Gets the mean across real (non-padding) tokens
    return average 


def get_embeddings(sentences, tok, model, layers, delimiter_token=None): # Function to get the embeddigns  
    inputs = tok(sentences, return_tensors="pt", padding=True, truncation=True, add_special_tokens=True) # Tokenizes my string data  
    for k in inputs: # Moves all tokenized tensors to the same device as the model
        inputs[k] = inputs[k].to(device) 

    start_indices = None # We don't use any special start positions unless there is a delimiter
    if delimiter_token is not None: # Only computes start positions if there is a delimiter (the EOS)
        delim_ids = tok(delimiter_token, add_special_tokens=False)["input_ids"] # Tokenizes the delimiter string without adding special tokens. The result is the raw token ID sequence for the delimiter 
        if delim_ids: # Proceeds only if there is a token ID (the EOS)
            anchor = delim_ids[-1] # Uses the last token id of the delimiter sequence as the anchor
            ids = inputs["input_ids"] 
            am  = inputs["attention_mask"] # Grabs the encoded batch and the mask (1=real token, 0=pad)
            B, T = ids.shape # Batch size: B, sequence length: T
            starts = [] # Collects the per-sequence start index 
            for b in range(B): # Goes over each sequence 
                real_T = int(am[b].sum().item()) # Counts how many real tokens the sequence has
                row = ids[b, :real_T].tolist() # Slices out only the real tokens and converts them into a list 
                occ = [i for i, t in enumerate(row) if t == anchor] # Finds all positions i where the "anchor" token id appears 
                internal = [i for i in occ if 0 < i < real_T - 1] # Chooses the first internal delimiter (only the ones that occur inside the real content). It filters out boundary hits (position 0, last real position real_T-1)
                if internal: # If there is an internal delimiter, start pooling after it
                    starts.append(internal[0] + 1)
                else:
                    starts.append(0)  # No internal delimiter here, so pool from start (no priming)
            start_indices = torch.tensor(starts, dtype=torch.long, device=device) # Converts the list into a tensor on the device


    with torch.no_grad(): # Runs a forward pass without tracking gradients to save memory
        out = model(**inputs, output_hidden_states=True, use_cache=False) # output_hidden_states=True asks the model to return the hidden states for every layer and use_cache=False avoids returning past_key_values 

    results = {} 
    for i in layers: # Goes over each layer
        layer_hidden = out.hidden_states[i] # Gets the hidden states at that layer
        pooled = mean_pool(layer_hidden, inputs["attention_mask"], start_indices=start_indices) # Masked mean
        results[i] = pooled.cpu().numpy() # Moves to CPU and NumPy
    return results 

# Calling the functions for each language:
embs_nt = get_embeddings(neutral, tokenizer, model, layers) # There is no priming, so when I don't use delimiter_token=tokenizer.eos_token, the whole sentence will be turned into vectors
print("Neutral OK")
embs_pt = get_embeddings(pt, tokenizer, model, layers, delimiter_token=tokenizer.eos_token)
print("Portugues OK")
embs_sp = get_embeddings(sp, tokenizer, model, layers, delimiter_token=tokenizer.eos_token)
print("Spanish OK")
embs_fr = get_embeddings(fr, tokenizer, model, layers, delimiter_token=tokenizer.eos_token)
print("French OK")

for i in layers: # Goes over layers to check if the shapes are correct
    print(f"Layer {i}: Neutral={embs_nt[i].shape}, PT={embs_pt[i].shape}, SP={embs_sp[i].shape}, FR={embs_fr[i].shape}")


def save_layer_parquet(layer_idx, nt_arr, pt_arr, sp_arr, fr_arr, outdir): # Function to save them in parquet as well
    nt_arr = nt_arr.astype(np.float32) # Ensures all input arrays are float32 for all languages
    pt_arr = pt_arr.astype(np.float32)
    sp_arr = sp_arr.astype(np.float32)
    fr_arr = fr_arr.astype(np.float32)

    n1, h1 = nt_arr.shape # Grabs (num_rows, hidden_dim) for each language/condition
    n2, h2 = pt_arr.shape
    n3, h3 = sp_arr.shape
    n4, h4 = fr_arr.shape
    assert n1 == n2 == n3 == n4 and h1 == h2 == h3 == h4 # Check if all of them have the same shape
    N, Hdim = n1, h1 # Makes sure they all have the same names (N = number of items/sentences, Hdim = embedding size)

    nt_flat = pa.array(nt_arr.reshape(-1), type=pa.float32()) # Flattens each (N, Hdim) to a 1D vector, wraps as arrow arrays for each language
    pt_flat = pa.array(pt_arr.reshape(-1), type=pa.float32())
    sp_flat = pa.array(sp_arr.reshape(-1), type=pa.float32())
    fr_flat = pa.array(fr_arr.reshape(-1), type=pa.float32())

    nt_list = pa.FixedSizeListArray.from_arrays(nt_flat, Hdim) # Converts the flat arrays into FixedSizeListArray for each language
    pt_list = pa.FixedSizeListArray.from_arrays(pt_flat, Hdim)
    sp_list = pa.FixedSizeListArray.from_arrays(sp_flat, Hdim)
    fr_list = pa.FixedSizeListArray.from_arrays(fr_flat, Hdim)
    # Builds an arrow table with one column per language
    table = pa.table({f"Neutral_L{layer_idx}": nt_list, f"Portuguese_L{layer_idx}": pt_list, f"Spanish_L{layer_idx}": sp_list, f"French_L{layer_idx}": fr_list})

    out_path = os.path.join(outdir, f"embeddings_L{layer_idx}_with_priming.parquet") # Outputs
    pq.write_table(table, out_path, compression="zstd") # Writing
    print(f"{out_path} (rows={N}, H={Hdim})") 

for layer in layers: # For each layer we are using, write the results
    save_layer_parquet(layer, embs_nt[layer], embs_pt[layer], embs_sp[layer], embs_fr[layer], outdir)

print("All parquet files saved.")