Finite Tractus: Code Companion

Core Code (JPEG Embedding Perturbation)

From Finite Tractus: The Hidden Geometry of Language and Thought (Haylett, 2025).

# -*- coding: utf-8 -*-
    """
    Created Feb 2025
    For Finite Tractus - JPEG Embedding Encoding
    @author: Kevin R. Haylett
    """
    import torch
    import torch.nn as nn
    from transformers import GPT2LMHeadModel, GPT2Tokenizer
    import numpy as np
    from PIL import Image
    import io
    from scipy.spatial.distance import cosine
    # --- Custom JPEG Compression Layer ---
    class JPEGCompressionLayer(nn.Module):
        def __init__(self, quality=100):
            super().__init__()
            self.quality = quality  # JPEG quality parameter

        def forward(self, embeddings):
            """
            embeddings: Tensor of shape [batch_size, seq_length, embedding_dim]
            Applies JPEG compression/decompression to each token embedding.
            """
            # Detach the tensor and convert to numpy for processing
            embeddings_np = embeddings.detach().cpu().numpy()
            batch_size, seq_length, embedding_dim = embeddings_np.shape
            
            # Process each token embedding individually
            processed_embeddings = np.empty_like(embeddings_np)
            for b in range(batch_size):
                for i in range(seq_length):
                    vec = embeddings_np[b, i, :]
                    processed_vec = self.jpeg_process(vec, quality=self.quality)
                    processed_embeddings[b, i, :] = processed_vec

            # Convert back to a torch tensor on the original device and with original dtype
            return torch.tensor(processed_embeddings, device=embeddings.device, dtype=embeddings.dtype)

        def jpeg_process(self, embedding, quality=95):
            """
            Process a 1D embedding vector:
              1. Ensure even length (pad if necessary).
              2. Reshape into a 2D array (2 rows).
              3. Normalize to [0, 255].
              4. Save as JPEG (simulate compression).
              5. Load and inverse normalize.
              6. Flatten back to 1D.
            """
            original_length = len(embedding)
            if original_length % 2 != 0:
                embedding = np.append(embedding, 0)  # pad to even length

            # Reshape into 2 rows
            reshaped = np.reshape(embedding, (2, -1))
            
            # Normalize to the 0-255 range
            min_val = reshaped.min()
            max_val = reshaped.max()
            norm = (reshaped - min_val) / (max_val - min_val + 1e-8) * 255.0
            norm_img = norm.astype(np.uint8)
            
            # Save the normalized image into an in-memory JPEG buffer
            buffer = io.BytesIO()
            image = Image.fromarray(norm_img)
            image.save(buffer, format='JPEG', quality=quality)
            buffer.seek(0)
            
            # Read back the JPEG image from the buffer
            decompressed_img = Image.open(buffer)
            decompressed_array = np.array(decompressed_img)
            
            # Inverse normalization
            decompressed = decompressed_array.astype(np.float32) / 255.0 * (max_val - min_val) + min_val
            
            # Flatten back to 1D and remove any padding
            processed_vec = decompressed.flatten()[:original_length]
            return processed_vec

    # --- Modified GPT2 Model that Inserts JPEG Compression into the Pipeline ---
    class ModifiedGPT2Model(GPT2LMHeadModel):
        def __init__(self, config):
            super().__init__(config)
           
            
           
            # Insert our custom JPEG compression layer
            self.jpeg_layer = JPEGCompressionLayer(quality=85)   #<<<<<<<< ENTER QUALITY HERE 1-100!

    def forward(self, input_ids=None, **kwargs):
        # If input_ids are provided, compute embeddings from them.
        if input_ids is not None:
            inputs_embeds = self.transformer.wte(input_ids)
        else:
            # Otherwise, expect inputs_embeds to be provided directly.
            inputs_embeds = kwargs.pop("inputs_embeds", None)
            if inputs_embeds is None:
                raise ValueError("Either input_ids or inputs_embeds must be provided.")

        # Apply the custom JPEG compression layer
        processed_embeds = self.jpeg_layer(inputs_embeds)
        
        # Remove both 'input_ids' and 'inputs_embeds' from kwargs to avoid duplication.
        kwargs.pop("input_ids", None)
        kwargs.pop("inputs_embeds", None)

        # Pass the processed embeddings to the transformer
        transformer_outputs = self.transformer(inputs_embeds=processed_embeds, **kwargs)
        hidden_states = transformer_outputs[0]
        logits = self.lm_head(hidden_states)
        return logits

    # --- Main function to run the model ---
    def main():
        # Load tokenizer and modified model
        tokenizer = GPT2Tokenizer.from_pretrained("gpt2-large")
        model = ModifiedGPT2Model.from_pretrained("gpt2-large")
        
        
        
        
        # Prepare an input text prompt using the tokenizer with return_attention_mask=True
        prompt ="Explain the concept of scientific method."
     
        
        
        encoded_input = tokenizer(prompt, return_tensors="pt")
        input_ids = encoded_input["input_ids"]
        attention_mask = encoded_input["attention_mask"]
        
        
        # --- Debug: Inspect embeddings before and after JPEG compression ---
        # Get original embeddings from the embedding layer
        original_embeds = model.transformer.wte(input_ids)
        # Process embeddings using the JPEG layer
        processed_embeds = model.jpeg_layer(original_embeds)
        
        # Flatten the embeddings and compute cosine similarity (average over batch)
        original_flat = original_embeds.flatten().detach().cpu().numpy()
        processed_flat = processed_embeds.flatten().detach().cpu().numpy()
        sim = 1 - cosine(original_flat, processed_flat)
        print("Cosine similarity between original and processed embeddings:", sim)
        
            
        
        # Generate output with adjusted parameters, passing attention_mask along with input_ids
        output = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=500,#50,
            do_sample=True,
            temperature=0.8,
            top_k=50,
            top_p=0.95
        )
        
        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
        
        print("Prompt:", prompt)
        print("Generated Output:", generated_text)

    if __name__ == "__main__":
        main()