Upload PISCO

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,18 @@
+{
+  "_name_or_path": "/scratch/1/user/mlouis/calmar/pisco_hub_models/mistral_with_mistral_labels",
+  "architectures": [
+    "PISCO"
+  ],
+  "auto_map": {
+    "AutoConfig": "modelling_pisco.PISCOConfig",
+    "AutoModel": "modelling_pisco.PISCO"
+  },
+  "compr_rate": 16,
+  "decoder_model_name": "mistralai/Mistral-7B-Instruct-v0.2",
+  "device_map": null,
+  "lora_r": 16,
+  "model_type": "PISCO",
+  "sep": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.44.2"
+}

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "top_p": null,
+  "transformers_version": "4.44.2"
+}

model-00001-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9791259c3ee866c55b5339d42fae3e67138e73b83148e1ac9746fe9c99cb1af
+size 4999790024

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75a2dd6edc2d77e2eccd979954f662b1221218e35ec2cdbc4cbb3f52d5887c37
+size 4974831120

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27dbfb30c4431410f6c65e6327e06373238f54fc58ff3541410506fdd44f9be1
+size 4676964024

modelling_pisco.py

@@ -0,0 +1,324 @@
+import warnings
+import os
+import torch
+from peft import LoraConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PretrainedConfig, AutoConfig, GenerationConfig
+from jinja2.exceptions import TemplateError
+
+
+def add_memory_tokens_to_inputs(input_ids: torch.Tensor, attention_mask: torch.Tensor, n_mem_tokens: int, tokenizer):
+    """
+    Concatenate the input ids with n_mem_tokens mem_tokens and update the corresponding attention mask
+    """
+    assert len(tokenizer.mem_tokens) == n_mem_tokens, f"{len(tokenizer.mem_tokens)} VS {n_mem_tokens}"
+    mem_tokens = torch.stack([tokenizer.mem_token_ids_pt] * input_ids.size(0), 0)
+    assert len(mem_tokens.size()) == 2
+    assert len(mem_tokens) == input_ids.size(0)
+    assert len(mem_tokens[0]) == n_mem_tokens
+    #mem_tokens = torch.full((input_ids.size(0), n_mem_tokens), tokenizer.mem_token_id, dtype=torch.long)
+    input_ids = torch.cat([input_ids, mem_tokens], dim=1)
+    attention_mask = torch.cat([attention_mask, torch.ones(input_ids.size(0), n_mem_tokens)], dim=1)
+    return input_ids, attention_mask
+
+
+class PISCOConfig(PretrainedConfig):
+
+    model_type = "PISCO"
+    def __init__(self,
+                decoder_model_name: str = "meta-llama/Llama-2-7b-chat-hf",
+                compr_rate: int = 16,
+                **kwargs):
+        super().__init__(**kwargs)
+
+        self.decoder_model_name = decoder_model_name # model name of decoder
+        self.compr_rate = compr_rate # compression rate
+        self.lora_r = 16
+        self.sep = True
+        
+        
+class PISCO(PreTrainedModel):
+    config_class = PISCOConfig
+    def __init__(self, cfg):
+        super().__init__(cfg)
+        self.decoder_model_name = cfg.decoder_model_name
+        self.sep = cfg.sep
+        self.compr_rate = cfg.compr_rate
+        
+        self.create_tokenizer(cfg)
+        
+        # Base model config but we modify vocab size since we added tokens (mainly the mem tokens)
+        decoder_config = AutoConfig.from_pretrained(cfg.decoder_model_name)
+        decoder_config.vocab_size = len(self.tokenizer)
+        
+        # Initializing placeholder model:
+        self.decoder = AutoModelForCausalLM.from_config(decoder_config, 
+                                                        attn_implementation='flash_attention_2', 
+                                                        torch_dtype=torch.bfloat16)
+        
+        peft_config = self.get_peft_config(cfg)
+
+        self.adapter_keys = []
+        self.decoder.add_adapter(peft_config, 'decoder_adapter')
+        self.decoder.set_adapter('decoder_adapter')
+        self.adapter_keys.append('decoder_adapter')
+        self.decoder.add_adapter(peft_config, 'encoder_adapter')
+        self.adapter_keys.append('encoder_adapter')
+        
+        self.generation_config = GenerationConfig(do_sample=False, top_p=None)
+        
+        print('a')
+        # self.decoder = self.decoder.to('cuda')
+        print('b')
+        if torch.cuda.is_available():
+            print('c')
+            # self.decoder = self.decoder.to('cuda')
+            print('d')
+
+    def create_tokenizer(self, cfg):
+        self.tokenizer = AutoTokenizer.from_pretrained(cfg.decoder_model_name, use_fast=True, padding_side='left')
+        
+        n_mem_tokens = 128 // cfg.compr_rate
+        mem_tokens = ['<MEM' + str(i) + '>' for i in range(n_mem_tokens)]
+        self.tokenizer.add_special_tokens({'additional_special_tokens': mem_tokens + ['<AE>', '<ENC>', '<SEP>']}) 
+        self.tokenizer.mem_tokens = mem_tokens
+        
+        self.tokenizer.mem_token_ids = [self.tokenizer.convert_tokens_to_ids(elt) for elt in self.tokenizer.mem_tokens]
+        self.tokenizer.mem_token_ids_pt = torch.LongTensor(self.tokenizer.mem_token_ids) # required later on for operations on tensors
+        
+        self.tokenizer.ae_token = '<AE>' # token for autoencoding on decoder side
+        self.tokenizer.ae_token_id = self.tokenizer.convert_tokens_to_ids('<AE>')
+        self.tokenizer.enc_token = '<ENC>' # token for autoencoding on compressor side
+        self.tokenizer.sep_token = '<SEP>' # sep token between document
+        self.tokenizer.sep_token_id = self.tokenizer.convert_tokens_to_ids('<SEP>')
+
+        # if pad token exists then use pad token, othrwise bos token
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.bos_token_id
+
+    def set_all_adapters(self):
+        if len(self.adapter_keys) > 0:
+            self.decoder.set_adapter(self.adapter_keys)
+
+    def get_peft_config(self, cfg: PISCOConfig) -> LoraConfig:
+        """
+        Builds the peft config
+        """
+        return LoraConfig(task_type="CAUSAL_LM", r=cfg.lora_r, lora_alpha=2* cfg.lora_r, target_modules='all-linear', lora_dropout=0.1)
+            
+    def compress(self, enc_input_ids, enc_attention_mask):
+        return self.compr_decoder(enc_input_ids, enc_attention_mask)
+
+    def replace_emb(self, enc_input_ids, compressed_embs, dec_input_ids):
+        """
+        Compression logic (either with decoder or with dedicated compressor)
+        """
+        indices = range(0, enc_input_ids.size(0) + 1, self.generation_top_k)            
+        input_embeds = self.replace_embeddings(compressed_embs, dec_input_ids, indices)
+        return input_embeds
+
+    def compr_decoder(self, input_ids, attention_mask):
+        """
+        Compression using the decoder
+        """
+        assert input_ids.size() == attention_mask.size(), f"{input_ids.size()} vs {attention_mask.size()}"
+        
+        # Switch adapter if we are training two different ones:
+        if 'encoder_adapter' in self.adapter_keys:
+            self.decoder.set_adapter('encoder_adapter')
+            
+        print(self.decoder.device, input_ids.device, attention_mask.device)
+
+        emb = self.decoder(input_ids=input_ids,
+                           attention_mask=attention_mask,
+                           output_hidden_states=True).hidden_states[-1]
+        mask = torch.isin(input_ids, self.tokenizer.mem_token_ids_pt.to(input_ids.device))
+        return emb[mask].reshape(emb.size(0), -1, emb.size(-1))
+    
+    def prepare_encoder_inputs_to_decoder(self, texts, max_length, q_texts=None):
+        if q_texts is not None:
+            texts_to_encode = [self.tokenizer.enc_token + self.tokenizer.bos_token + '\nQuery:\n' + query + 'Document:\n' + text + self.tokenizer.eos_token 
+                               for text, query in zip(texts, q_texts)]
+            inp_enc = self.tokenizer(texts_to_encode, return_tensors='pt', padding='max_length', max_length=max_length + 8, truncation=True, add_special_tokens=False)
+        else:
+            inp_enc = [self.tokenizer.enc_token + self.tokenizer.bos_token + text + self.tokenizer.eos_token for text in texts]
+            inp_enc = self.tokenizer(inp_enc, return_tensors='pt', padding="longest", max_length=max_length+3, truncation=True, add_special_tokens=False)
+            
+        num_mem_tokens = 128 // self.compr_rate  # maybe change that
+        assert num_mem_tokens == len(self.tokenizer.mem_tokens)
+        inp_enc['input_ids'], inp_enc['attention_mask'] = add_memory_tokens_to_inputs(inp_enc['input_ids'], 
+                                                                                        inp_enc['attention_mask'], 
+                                                                                        num_mem_tokens, 
+                                                                                        tokenizer=self.tokenizer)
+        
+        return inp_enc
+    
+    def prepare_encoder_inputs(self, texts, max_length):
+        return self.prepare_encoder_inputs_to_decoder(texts, max_length)
+        
+    def replace_embeddings(self, compressed_embs, dec_input_ids, indices):
+        """
+        Replace memory tokens in the decoder input to with the compressed embeddings
+        """
+        inputs_embeds = self.decoder.get_input_embeddings()(dec_input_ids)
+        num_embs = compressed_embs.size(1)
+        if self.sep:
+            slot_len = num_embs + 1
+        else:
+            slot_len = num_embs
+        # get first mem_token indices
+        first_mem_token_indices = torch.argmax((dec_input_ids == self.tokenizer.mem_token_ids[0]).int(), dim=1)
+        batch_size = inputs_embeds.size(0)
+        # for each example in batch, replace them with compressed embeddings
+        for i in range(batch_size):
+            for j in range(indices[i], indices[i + 1]):
+                start_idx = first_mem_token_indices[i].item() + (j-indices[i]) * slot_len
+                assert inputs_embeds[i, start_idx:start_idx + num_embs, :].size() == compressed_embs[j].size(), \
+                    f"{inputs_embeds[i, start_idx:start_idx + num_embs, :].size()} VS {compressed_embs[j].size()}"
+                inputs_embeds[i, start_idx:start_idx + num_embs, :] = compressed_embs[j]
+        return inputs_embeds
+
+    def forward(self,
+                enc_input_ids: torch.LongTensor = None,
+                enc_attention_mask: torch.LongTensor = None,
+                dec_input_ids: torch.LongTensor = None,
+                dec_attention_mask: torch.LongTensor = None,
+                labels: torch.LongTensor = None):
+        """
+        enc_input_ids: stores the contexts, should be flattened from all queries before input, can be of shape:
+            - (batch_size*generation_top_k, enc_token_length)
+            - (batch_size, generation_top_k, enc_token_length)
+        enc_attention_mask: attention mask of enc_input_ids, same shape as enc_input_ids
+        dec_input_ids: stores the prompts (including mem tokens), dimention (batch_size, dec_token_length)
+        dec_attention_mask: attention mask of dec_input_ids
+        """ 
+        assert enc_input_ids.size() == enc_attention_mask.size(), f"{enc_input_ids.size()} vs {enc_attention_mask.size()}"
+        
+        if len(enc_input_ids.size()) == 3: # likely from bergen: we just flatten all of this to perform encoding in one batch
+            batch_size, top_k, seq_length = enc_input_ids.size()
+            enc_input_ids = enc_input_ids.view(batch_size * top_k, seq_length)
+            enc_attention_mask = enc_attention_mask.view(batch_size * top_k, seq_length)
+        
+        # Here, we should have top_k times more elements in enc_input_ids than in dec_input_ids
+        assert enc_input_ids.size(0) == dec_input_ids.size(0) * self.generation_top_k, \
+            f"{enc_input_ids.size(0)} VS {dec_input_ids.size(0)} with generation_top_k={self.generation_top_k}"
+            
+        # Perform compression with gradient tracking
+        compressed_embs = self.compress(enc_input_ids, enc_attention_mask)
+        inputs_embeds = self.replace_emb(enc_input_ids, compressed_embs, dec_input_ids)
+
+        # decoding
+        if 'decoder_adapter' in self.adapter_keys:
+            self.decoder.set_adapter('decoder_adapter')
+
+        decoder_outputs = self.decoder(inputs_embeds=inputs_embeds, attention_mask=dec_attention_mask, labels=labels)
+
+        # At end of forward, we need to activate all adapters so that they are both trained...
+        self.set_all_adapters()
+
+        return {"loss": decoder_outputs.loss, "logits": decoder_outputs.logits}
+    
+    def generate_from_text(self, questions: list[str], documents: list[list[str]], max_new_tokens: int = 128) -> list[str]:
+        # TODO: test
+        self.generation_top_k = len(documents[0])
+        assert len(documents) == len(questions)
+        assert all([len(context) == len(documents[0]) for context in documents])
+        flat_documents = sum(documents, [])
+        
+        model_input = {}
+        input_encoder = self.prepare_encoder_inputs(flat_documents, max_length=128)
+        device = self.decoder.device
+        
+        model_input['enc_input_ids'], model_input['enc_attention_mask'] = input_encoder['input_ids'].to(device), input_encoder['attention_mask'].to(device)
+        
+        instr = [self.blend_prompt_and_memory_tokens(query=q) for q in questions]
+
+        inp_dec = self.tokenizer(instr, return_tensors='pt', padding="longest", add_special_tokens=False, truncation=True,  max_length=2048)
+        
+        model_input['dec_input_ids'], model_input['dec_attention_mask'] = inp_dec['input_ids'].to(device), inp_dec['attention_mask'].to(device)
+        
+        return self.generate(model_input, max_new_tokens=max_new_tokens)
+        
+    def compress_documents(self, documents: list[str]) -> torch.Tensor:
+        # TODO: test
+        input_encoder = self.prepare_encoder_inputs(documents, max_length=128)
+        enc_input_ids = input_encoder['input_ids'].to(self.decoder.device)
+        attention_mask = input_encoder['attention_mask'].to(self.decoder.device)
+        print('yo', self.decoder.device, enc_input_ids.device, attention_mask.device)
+        return self.compress(enc_input_ids=enc_input_ids, enc_attention_mask=attention_mask)
+
+    def generate(self, model_input, max_new_tokens=128):
+
+        enc_input_ids, enc_attention_mask, dec_input_ids, dec_attention_mask = model_input['enc_input_ids'], model_input['enc_attention_mask'], model_input['dec_input_ids'], model_input['dec_attention_mask']
+        
+        print('in gen')
+        print(enc_input_ids.size())
+        print(dec_input_ids.size())
+        
+        assert enc_input_ids.size() == enc_attention_mask.size()
+        
+        if len(enc_input_ids.size()) == 3: # likely from bergen: we just flatten all of this to perform encoding in one batch
+            batch_size, top_k, seq_length = enc_input_ids.size()
+            enc_input_ids = enc_input_ids.view(batch_size * top_k, seq_length)
+            enc_attention_mask = enc_attention_mask.view(batch_size * top_k, seq_length)
+            
+        # Here, we should have top_k times more elements in enc_input_ids than in dec_input_ids
+        assert enc_input_ids.size(0) == dec_input_ids.size(0) * self.generation_top_k, \
+            f"{enc_input_ids.size(0)} VS {dec_input_ids.size(0)} with generation_top_k={self.generation_top_k}"
+            
+        compressed_embs = self.compress(enc_input_ids, enc_attention_mask)
+        inputs_embeds = self.replace_emb(enc_input_ids, compressed_embs, dec_input_ids)
+        
+        # Switch adapter if we are training two different ones:
+        if 'decoder_adapter' in self.adapter_keys:
+            self.decoder.set_adapter('decoder_adapter')
+            
+
+        output_ids = self.decoder.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=dec_attention_mask,
+            generation_config=self.generation_config,
+            max_new_tokens=max_new_tokens
+            )
+
+        decoded = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+        
+        return decoded
+    
+    def blend_prompt_and_memory_tokens(self, query: str):
+        """
+        Takes care of blending the prompt with the memory tokens:
+        Also returns, if a label is provided, the position of the first token index of the label (for loss comp later on)
+        """        
+
+        mem_tokens_str = ''.join(self.tokenizer.mem_tokens)        
+        mem_tokens_str += self.tokenizer.sep_token
+        
+        # proper names for "eval" call, don't remove these lines
+        docs = mem_tokens_str * self.generation_top_k
+        question = query
+        
+        prompt_system = 'You are a helpful assistant. Your task is to extract relevant information from provided documents and to answer to questions as briefly as possible.'
+        prompt_user = f"Background:\n{docs}\n\nQuestion:{question}"
+        
+        # Prepare the messages with system and user roles
+        messages = [
+            {"role": "system", "content": prompt_system},
+            {"role": "user", "content": prompt_user.replace(':\ ', ': ')}
+        ]
+
+        # Attempt to apply the system role and catch if it's not supported
+        try:
+            prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            
+        except TemplateError as e:
+            # Catch the error related to system role and handle it (e.g. gemma)
+            if "System role not supported" in str(e):
+                # Remove system role and proceed with only the user role
+                messages = [{"role": "user", "content": messages[0]['content'] + '\n' + messages[1]['content']}]
+                # Apply template again without system role
+                prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            else:
+                # Re-raise the exception if it's unrelated to system role
+                raise e
+
+        return prompt