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 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('')[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" prompt = data.pop("inputs", None) parameters = data.pop("parameters", None) mode = data.pop('mode', 'diff2compound') 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