handler.py

from typing import Dict, List, Any
import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import PreTrainedTokenizerFast
from transformers import GenerationConfig
import transformers 
import pandas as pd
import time
import numpy as np
from precious3_gpt_multi_modal import Custom_MPTForCausalLM


class EndpointHandler:
    def __init__(self, path=""):

        self.path = path
        # load model and processor from path
        self.model = Custom_MPTForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16).to('cuda')
        self.tokenizer = PreTrainedTokenizerFast(tokenizer_file = os.path.join(path, "tokenizer.json"), unk_token="[UNK]",
            pad_token="[PAD]",
            eos_token="[EOS]",
            bos_token="[BOS]")
        self.model.config.pad_token_id = self.tokenizer.pad_token_id
        self.model.config.bos_token_id = self.tokenizer.bos_token_id
        self.model.config.eos_token_id = self.tokenizer.eos_token_id
        unique_entities_p3 = pd.read_csv(os.path.join(path, 'p3_entities_with_type.csv'))
        self.unique_compounds_p3 = [i.strip() for i in unique_entities_p3[unique_entities_p3.type=='compound'].entity.to_list()]
        
        self.emb_gpt_genes = pd.read_pickle(os.path.join(self.path, 'multi-modal-data/emb_gpt_genes.pickle'))
        self.emb_hgt_genes = pd.read_pickle(os.path.join(self.path, 'multi-modal-data/emb_hgt_genes.pickle'))


    def create_prompt(self, prompt_config):

        prompt = "[BOS]"

        multi_modal_prefix = '<modality0><modality1><modality2><modality3>'*3

        for k, v in prompt_config.items():
            if k=='instruction':
                prompt+=f"<{v}>"
            elif k=='up':
                prompt+=f'{multi_modal_prefix}<{k}>{v}</{k}>' if isinstance(v, str) else f'{multi_modal_prefix}<{k}>{" ".join(v)} </{k}>'
            elif k=='down':
                prompt+=f'{multi_modal_prefix}<{k}>{v}</{k}>' if isinstance(v, str) else f'{multi_modal_prefix}<{k}>{" ".join(v)} </{k}>'
            else:
                prompt+=f'<{k}>{v}</{k}>' if isinstance(v, str) else f'<{k}>{" ".join(v)} </{k}>'
        return prompt

    def custom_generate(self,
                        input_ids, 
                        acc_embs_up_kg_mean, 
                        acc_embs_down_kg_mean, 
                        acc_embs_up_txt_mean, 
                        acc_embs_down_txt_mean,
                        device, 
                        max_new_tokens,
                        unique_compounds,
                        temperature=0.8, 
                        top_p=0.2, top_k=3550, 
                        n_next_tokens=50, num_return_sequences=1):
        torch.manual_seed(137)

        # Set parameters
        # temperature - Higher value for more randomness, lower for more control
        # top_p - Probability threshold for nucleus sampling (aka top-p sampling)
        # top_k - Ignore logits below the top-k value to reduce randomness (if non-zero)
        # n_next_tokens - Number of top next tokens when predicting compounds

        modality0_emb =  torch.unsqueeze(torch.from_numpy(acc_embs_up_kg_mean), 0).to(device) 
        modality1_emb =  torch.unsqueeze(torch.from_numpy(acc_embs_down_kg_mean), 0).to(device)
        modality2_emb =  torch.unsqueeze(torch.from_numpy(acc_embs_up_txt_mean), 0).to(device)
        modality3_emb =  torch.unsqueeze(torch.from_numpy(acc_embs_down_txt_mean), 0).to(device)


        # Generate sequences
        outputs = []
        next_token_compounds = [] 

        for _ in range(num_return_sequences):
            start_time = time.time()
            generated_sequence = []
            current_token = input_ids.clone()

            for _ in range(max_new_tokens):  # Maximum length of generated sequence
                # Forward pass through the model
                logits = self.model.forward(input_ids=current_token, 
                                           modality0_emb=modality0_emb, # torch.tensor(efo_embeddings['EFO_0002618'], dtype=torch.bfloat16).to(device), 
                                           modality0_token_id=62191, 
                                           modality1_emb=modality1_emb, # torch.tensor(efo_embeddings['EFO_0002618'], dtype=torch.bfloat16).to(device), 
                                           modality1_token_id=62192,
                                           modality2_emb=modality2_emb, # torch.tensor(efo_embeddings['EFO_0002618'], dtype=torch.bfloat16).to(device), 
                                           modality2_token_id=62193, 
                                           modality3_emb=modality3_emb, # torch.tensor(efo_embeddings['EFO_0002618'], dtype=torch.bfloat16).to(device), 
                                           modality3_token_id=62194)[0]

                # Apply temperature to logits
                if temperature != 1.0:
                    logits = logits / temperature

                # Apply top-p sampling (nucleus sampling)
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p

                if top_k > 0:
                    sorted_indices_to_remove[..., top_k:] = 1

                # Set the logit values of the removed indices to a very small negative value
                inf_tensor = torch.tensor(float("-inf")).type(torch.bfloat16).to(logits.device)

                logits = logits.where(sorted_indices_to_remove, inf_tensor)


                # Sample the next token
                if current_token[0][-1] == self.tokenizer.encode('<drug>')[0]:
                    next_token_compounds.append(torch.topk(torch.softmax(logits, dim=-1)[0][len(current_token[0])-1, :].flatten(), n_next_tokens).indices)

                next_token = torch.multinomial(torch.softmax(logits, dim=-1)[0], num_samples=1)[len(current_token[0])-1, :].unsqueeze(0)


                # Append the sampled token to the generated sequence
                generated_sequence.append(next_token.item())

                # Stop generation if an end token is generated
                if next_token == self.tokenizer.eos_token_id:
                    break

                # Prepare input for the next iteration
                current_token = torch.cat((current_token, next_token), dim=-1)
            print(time.time()-start_time)
            outputs.append(generated_sequence)

        predicted_compounds_ids = [self.tokenizer.convert_ids_to_tokens(j) for j in next_token_compounds]
        predicted_compounds = []
        for j in predicted_compounds_ids:
            predicted_compounds.append([i.strip() for i in j])
        return outputs, predicted_compounds


    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
        """
        Args:
            data (:dict:):
                The payload with the text prompt and generation parameters.
        """
        torch.manual_seed(137)
        
        device = "cuda"
        config_data = data.pop("inputs", None)
        parameters = data.pop("parameters", None)
        mode = data.pop('mode', 'diff2compound')
        
        prompt = self.create_prompt(config_data)
        
        inputs = self.tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to(device)

        max_new_tokens = 600 - len(input_ids[0]) 
        try:
            acc_embs_up1 = []
            acc_embs_up2 = []
            for gs in config_data['up']: 
                try:
                    acc_embs_up1.append(self.emb_hgt_genes[self.emb_hgt_genes.gene_symbol==gs].embs.values[0])
                    acc_embs_up2.append(self.emb_gpt_genes[self.emb_gpt_genes.gene_symbol==gs].embs.values[0])
                except Exception as e: 
                    pass
            acc_embs_up1_mean = np.array(acc_embs_up1).mean(0)
            acc_embs_up2_mean = np.array(acc_embs_up2).mean(0)

            acc_embs_down1 = []
            acc_embs_down2 = []
            for gs in config_data['down']:
                try:
                    acc_embs_down1.append(self.emb_hgt_genes[self.emb_hgt_genes.gene_symbol==gs].embs.values[0])
                    acc_embs_down2.append(self.emb_gpt_genes[self.emb_gpt_genes.gene_symbol==gs].embs.values[0])
                except Exception as e: 
                    pass
            acc_embs_down1_mean = np.array(acc_embs_down1).mean(0)
            acc_embs_down2_mean = np.array(acc_embs_down2).mean(0)

            generated_sequence, raw_next_token_generation = self.custom_generate(input_ids = input_ids, 
                                                             acc_embs_up_kg_mean=acc_embs_up1_mean,
                                                             acc_embs_down_kg_mean=acc_embs_down1_mean, 
                                                             acc_embs_up_txt_mean=acc_embs_up2_mean,
                                                             acc_embs_down_txt_mean=acc_embs_down2_mean, max_new_tokens=max_new_tokens,
                                                             device=device, unique_compounds=self.unique_compounds_p3, **parameters)
            next_token_generation = [sorted(set(i) & set(self.unique_compounds_p3), key = i.index) for i in raw_next_token_generation]

            if mode == "meta2diff":
                outputs = {"up": generated_sequence, "down": generated_sequence}
            else:
                outputs = generated_sequence
            out = {"output": outputs, 'compounds': next_token_generation, "raw_output": raw_next_token_generation, "mode": mode, 'message': "Done!"}
            
        except Exception as e:
            print(e)
            outputs, next_token_generation = [None], [None]
            out = {"output": outputs, "mode": mode, 'message': f"{e}"}

        return out