Delete g_retriever

g_retriever/g_retriever_config.py

@@ -1,31 +0,0 @@
-from transformers import LlamaConfig
-
-class GRetrieverConfig(LlamaConfig):
-    model_type = "llama"
-    
-    def __init__(
-        self,
-        max_txt_len: int = 1024,
-        max_new_tokens: int = 256,
-        gnn_num_layers: int = 4,
-        gnn_in_dim: int = 768,
-        gnn_hidden_dim: int = 1024,
-        gnn_num_heads: int = 4,
-        gnn_dropout: int = 0,
-        bos_id: list = [128000, 128006, 882, 128007],
-        **kwargs
-    ):
-        pretrained_config = LlamaConfig.from_pretrained("NousResearch/Hermes-3-Llama-3.1-8B")
-        pretrained_config.update(kwargs)
-        
-        self.max_txt_len = max_txt_len
-        self.max_new_tokens = max_new_tokens
-        self.gnn_num_layers = gnn_num_layers
-        self.gnn_in_dim = gnn_in_dim
-        self.gnn_hidden_dim = gnn_hidden_dim
-        self.gnn_num_heads = gnn_num_heads
-        self.gnn_dropout = gnn_dropout
-        self.bos_id = bos_id
-        
-        super().__init__(**pretrained_config.to_dict())
-        self.pad_token_id = pretrained_config.eos_token_id

g_retriever/g_retriever_model.py

@@ -1,215 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-from transformers import LlamaForCausalLM
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.cache_utils import StaticCache
-from transformers.models.llama.modeling_llama import _prepare_4d_causal_attention_mask_with_cache_position
-from .g_retriever_config import GRetrieverConfig
-from .gnn import GAT
-
-from functools import wraps
-from torch_geometric.nn.pool import global_mean_pool
-
-class GRetrieverModel(LlamaForCausalLM):
-    config_class = GRetrieverConfig
-    
-    def __init__(self, config):
-        super().__init__(config)
-        self.graph_encoder = GAT(
-            in_channels=config.gnn_in_dim,
-            out_channels=config.gnn_hidden_dim,
-            hidden_channels=config.gnn_hidden_dim,
-            num_layers=config.gnn_num_layers,
-            dropout=config.gnn_dropout,
-            num_heads=config.gnn_num_heads,
-        ).to(self.model.dtype)
-
-        self.projector = nn.Sequential(
-            nn.Linear(config.gnn_hidden_dim, 2048),
-            nn.Sigmoid(),
-            nn.Linear(2048, self.get_input_embeddings().embedding_dim),
-        ).to(self.model.dtype)
-    
-    def encode_graphs(self, graph):
-        n_embeds, _ = self.graph_encoder(
-            graph.x.to(self.model.dtype),
-            graph.edge_index.long(),
-            graph.edge_attr.to(self.model.dtype)
-        )
-
-        # mean pooling
-        g_embeds = global_mean_pool(n_embeds, graph.batch.to(n_embeds.device))
-
-        return g_embeds
-    
-    @wraps(LlamaForCausalLM.forward)
-    def forward(
-        self,
-        input_ids=None,
-        graph=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        return_dict=None,
-        cache_position=None
-    ):  
-        inputs = input_ids.clone()
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if (inputs==-1).any():
-            # embed bos prompt
-            bos_embeds = self.get_input_embeddings()(torch.tensor(
-                self.config.bos_id,
-                device=self.model.device
-            ))
-            
-            # encode graph
-            graph_embeds = self.encode_graphs(graph)
-            graph_embeds = self.projector(graph_embeds).to(self.model.device)
-            
-            # prepare for reserved ids (bos+graph)
-            non_tokenized_ids = (inputs == -1).nonzero()
-            non_tokenized_shape = non_tokenized_ids[:, 0], non_tokenized_ids[:, 1]
-            
-            # embed inputs
-            inputs[non_tokenized_shape] = self.config.pad_token_id
-            temp_inputs_embeds = self.get_input_embeddings()(inputs)
-            non_tokenized_embeds = torch.cat([bos_embeds.repeat(len(inputs), 1, 1), graph_embeds.unsqueeze(1)], dim=1)
-            
-            # replace reserved ids with bos+graph
-            inputs_embeds = temp_inputs_embeds.clone()
-            inputs_embeds[non_tokenized_shape] = non_tokenized_embeds.view(len(non_tokenized_ids), -1)
-        
-        else:
-            inputs_embeds = self.get_input_embeddings()(inputs)
-        
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs = self.model(
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            cache_position=cache_position,
-        )
-
-        hidden_states = outputs[0]
-        
-        if self.config.pretraining_tp > 1:
-            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
-            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
-            logits = torch.cat(logits, dim=-1)
-        else:
-            logits = self.lm_head(hidden_states)
-        logits = logits.float()
-
-        loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
-            loss_fct = nn.CrossEntropyLoss()
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
-            # Enable model parallelism
-            shift_labels = shift_labels.to(shift_logits.device)
-            loss = loss_fct(shift_logits, shift_labels)
-
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
-        
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-    
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        graph=None,
-        past_key_values=None,
-        attention_mask=None,
-        inputs_embeds=None,
-        cache_position=None,
-        position_ids=None,
-        use_cache=True,
-        **kwargs,
-    ):
-        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
-        # Exception 1: when passing input_embeds, input_ids may be missing entries
-        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
-        if past_key_values is not None:
-            if inputs_embeds is not None:  # Exception 1
-                input_ids = input_ids[:, -cache_position.shape[0] :]
-            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
-                input_ids = input_ids[:, cache_position]
-
-        if attention_mask is not None and position_ids is None:
-            # create position_ids on the fly for batch generation
-            position_ids = attention_mask.long().cumsum(-1) - 1
-            position_ids.masked_fill_(attention_mask == 0, 1)
-            if past_key_values:
-                position_ids = position_ids[:, -input_ids.shape[1] :]
-
-                # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s  `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture.
-                position_ids = position_ids.clone(memory_format=torch.contiguous_format)
-
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and cache_position[0] == 0:
-            model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
-        else:
-            # The clone here is for the same reason as for `position_ids`.
-            model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None}
-
-        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
-            if model_inputs["inputs_embeds"] is not None:
-                batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
-                device = model_inputs["inputs_embeds"].device
-            else:
-                batch_size, sequence_length = model_inputs["input_ids"].shape
-                device = model_inputs["input_ids"].device
-
-            dtype = self.lm_head.weight.dtype
-            min_dtype = torch.finfo(dtype).min
-
-            attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(
-                attention_mask,
-                sequence_length=sequence_length,
-                target_length=past_key_values.get_max_length(),
-                dtype=dtype,
-                device=device,
-                min_dtype=min_dtype,
-                cache_position=cache_position,
-                batch_size=batch_size,
-            )
-
-        model_inputs.update(
-            {
-                "graph": graph,
-                "position_ids": position_ids,
-                "cache_position": cache_position,
-                "past_key_values": past_key_values,
-                "use_cache": use_cache,
-                "attention_mask": attention_mask,
-            }
-        )
-        return model_inputs

g_retriever/g_retriever_pipeline.py

@@ -1,51 +0,0 @@
-from transformers import Pipeline, AutoTokenizer
-from torch_geometric.data import Batch
-import torch
-
-class GRetrieverPipeline(Pipeline):
-    def __init__(self, **kwargs):
-        Pipeline.__init__(self, **kwargs)
-        
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model.config._name_or_path)
-        self.eos_user = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>"
-        self.max_txt_len = self.model.config.max_txt_len
-        self.bos_length = len(self.model.config.bos_id)
-        
-    def _sanitize_parameters(self, **kwargs):
-        preprocess_kwargs = {}
-        if "textualized_graph" in kwargs:
-            preprocess_kwargs["textualized_graph"] = kwargs["textualized_graph"]
-
-        if "graph" in kwargs:
-            preprocess_kwargs["graph"] = kwargs["graph"]
-            
-        return preprocess_kwargs, {}, {}
-
-    def preprocess(self, inputs, textualized_graph, graph):
-        textualized_graph_ids = self.tokenizer(textualized_graph, add_special_tokens=False)["input_ids"][:self.max_txt_len]
-        question_ids = self.tokenizer(inputs, add_special_tokens=False)["input_ids"]
-        eos_user_ids = self.tokenizer(self.eos_user, add_special_tokens=False)["input_ids"]
-        
-        input_ids = torch.tensor([
-            [-1]*(self.bos_length + 1)
-            + textualized_graph_ids
-            + question_ids
-            + eos_user_ids
-        ])
-        model_inputs = {
-            "input_ids": input_ids,
-            "attention_mask": torch.ones_like(input_ids)
-        }
-        model_inputs.update({
-            "graph": Batch.from_data_list([graph])
-        })
-        
-        return model_inputs
-
-    def _forward(self, model_inputs):
-        model_outputs = self.model.generate(**model_inputs)
-            
-        return model_outputs
-
-    def postprocess(self, model_outputs):
-        return model_outputs