Functions — github.com/OSU-NLP-Group/In-Context-Reranking

hub / github.com/OSU-NLP-Group/In-Context-Reranking / functions

Functions90 in github.com/OSU-NLP-Group/In-Context-Reranking

⨍ Functions 90 ◇ Types & classes 25 ↳ Endpoints 2

Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. Parameters: query_states (`tor

src/custom/custom_cache.py:14

↓ 9 callersMethod__init__

(self, config)

src/custom/custom_modeling_llama.py:296

↓ 7 callersMethod__init__

(self, config)

src/custom/custom_modeling_mistral.py:151

↓ 7 callersMethod_prepare_input_for_document_retrieval

Only tested with Mistral and Llama-3.1. Models using other tokenizers may need to modify this function.

src/in_context_reranker.py:426

↓ 7 callersFunctionrepeat_kv

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, he

src/custom/custom_modeling_mistral.py:165

↓ 7 callersMethodscore_documents

( self, llm_input, doc_tok_idx_spans, query_start_tok_idx,

src/in_context_reranker.py:209

↓ 4 callersFunctionrepeat_kv

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, he

src/custom/custom_modeling_llama.py:329

↓ 4 callersMethodrerank

Rerank the documents based on the query using a sliding window strategy. Assume that input documents are sorted by their relevance to

src/rank_gpt_reranker.py:176

↓ 3 callersMethod__show_tokens

(self, string)

src/in_context_reranker.py:536

↓ 3 callersFunctionapply_rotary_pos_emb

Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): T

src/custom/custom_modeling_llama.py:268

↓ 3 callersFunctionapply_rotary_pos_emb

Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): T

src/custom/custom_modeling_mistral.py:123

↓ 2 callersMethod_get_openai_ranking

(self, text: str, max_tokens)

src/rank_gpt_reranker.py:59

↓ 2 callersFunction_prepare_4d_causal_attention_mask_with_cache_position

Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`,

src/custom/custom_modeling_llama.py:73

↓ 2 callersMethodfrom_legacy_cache

Converts a cache in the legacy cache format into an equivalent `DynamicCache`.

src/custom/custom_cache.py:61

↓ 2 callersFunctionrotate_half

Rotates half the hidden dims of the input.

src/custom/custom_modeling_llama.py:261

↓ 2 callersFunctionrotate_half

Rotates half the hidden dims of the input.

src/custom/custom_modeling_mistral.py:115

↓ 1 callersMethod_create_prompt

(self, query, doc_pool)

src/rank_gpt_reranker.py:77

↓ 1 callersMethod_dynamic_frequency_update

dynamic RoPE layers should recompute `inv_freq` in the following situations: 1 - growing beyond the cached sequence length (allow sca

src/custom/custom_modeling_llama.py:194

↓ 1 callersMethod_get_attn_weights

(cls, key_states, query_states, use_cpu=False)

src/in_context_reranker.py:543

↓ 1 callersMethod_get_causal_mask

(cls, attn_weights)

src/in_context_reranker.py:574

↓ 1 callersMethod_rank_docs_from_output

(self, output, doc_pool)

src/rank_gpt_reranker.py:113

↓ 1 callersMethod_setup_llm_prompts

(self, prompt_template, base_llm_name)

src/in_context_reranker.py:128

↓ 1 callersMethod_update_causal_mask

( self, attention_mask: torch.Tensor, input_tensor: torch.Tensor, cache_positi

src/custom/custom_modeling_llama.py:1073

↓ 1 callersMethod_update_causal_mask

( self, attention_mask: torch.Tensor, input_tensor: torch.Tensor, cache_positi

src/custom/custom_modeling_mistral.py:868

↓ 1 callersFunctionbm25_retrieve_beir

(task, K)

src/bm25_retrieval.py:6

↓ 1 callersMethodforward

(self, x)

src/custom/custom_modeling_llama.py:306

↓ 1 callersMethodforward

(self, hidden_state)

src/custom/custom_modeling_mistral.py:160

↓ 1 callersMethodget_sorted_docs

Sort documents based on permutation method as seen in RankGPT

src/rank_gpt_reranker.py:162

↓ 1 callersMethodget_sorted_docs

(self, query, retrieval_doc_pool, return_per_doc_results=False, prompt_prefix='', order='desc')

src/in_context_reranker.py:304

Method__init__

(self, base_llm_name, tokenizer='', use_vllm=True,

src/rank_gpt_reranker.py:12

Method__init__

Inputs: base_llm: The base LLM model to be used for document retrieval. tokenizer: The tokenizer for the base LLM mod

src/in_context_reranker.py:16

Method__init__

LlamaRMSNorm is equivalent to T5LayerNorm

src/custom/custom_modeling_llama.py:127

Method__init__

( self, dim=None, max_position_embeddings=2048, base=10000, device=Non

src/custom/custom_modeling_llama.py:150

Method__init__

(self, *args, **kwargs)

src/custom/custom_modeling_llama.py:239

Method__init__

(self, *args, **kwargs)

src/custom/custom_modeling_llama.py:251

Method__init__

(self, config: LlamaConfig, layer_idx: Optional[int] = None)

src/custom/custom_modeling_llama.py:344

Method__init__

(self, *args, **kwargs)

src/custom/custom_modeling_llama.py:487

Method__init__

(self, config: LlamaConfig, layer_idx: int)

src/custom/custom_modeling_llama.py:719

Method__init__

(self, config: LlamaConfig)

src/custom/custom_modeling_llama.py:929

Method__init__

(self, config)

src/custom/custom_modeling_llama.py:1148

Method__init__

MistralRMSNorm is equivalent to T5LayerNorm

src/custom/custom_modeling_mistral.py:66

Method__init__

(self, dim, max_position_embeddings=2048, base=10000, device=None)

src/custom/custom_modeling_mistral.py:86

Method__init__

(self, config: MistralConfig, layer_idx: Optional[int] = None)

src/custom/custom_modeling_mistral.py:183

Method__init__

(self, *args, **kwargs)

src/custom/custom_modeling_mistral.py:289

Method__init__

(self, config: MistralConfig, layer_idx: int)

src/custom/custom_modeling_mistral.py:527

Method__init__

(self, config: MistralConfig)

src/custom/custom_modeling_mistral.py:727

Method__init__

(self, config)

src/custom/custom_modeling_mistral.py:980

Method__init__

(self, query_indices=[])

src/custom/custom_cache.py:9

Method_get_sorted_docs_from_prompts

Sort documents based on permutation method as seen in RankGPT

src/rank_gpt_reranker.py:90

Function_get_unpad_data

(attention_mask)

src/custom/custom_modeling_llama.py:62

Method_init_weights

(self, module)

src/custom/custom_modeling_llama.py:831

Method_init_weights

(self, module)

src/custom/custom_modeling_mistral.py:633

Methodextra_repr

(self)

src/custom/custom_modeling_llama.py:142

Methodextra_repr

(self)

src/custom/custom_modeling_mistral.py:81

Methodforward

(self, hidden_states)

src/custom/custom_modeling_llama.py:135

Methodforward

(self, x, position_ids)

src/custom/custom_modeling_llama.py:213

Methodforward

( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,

src/custom/custom_modeling_llama.py:379

Methodforward

( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.LongTensor] = None

src/custom/custom_modeling_llama.py:495

Methodforward

( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,

src/custom/custom_modeling_llama.py:620

Methodforward

Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (

src/custom/custom_modeling_llama.py:729

Methodforward

( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = N

src/custom/custom_modeling_llama.py:952

Methodforward

r""" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing

src/custom/custom_modeling_llama.py:1177

Methodforward

(self, hidden_states)

src/custom/custom_modeling_mistral.py:74

Methodforward

(self, x, position_ids)

src/custom/custom_modeling_mistral.py:98

Methodforward

( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,

src/custom/custom_modeling_mistral.py:215

Methodforward

( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,

src/custom/custom_modeling_mistral.py:297

Methodforward

( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,

src/custom/custom_modeling_mistral.py:437

Methodforward

Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (

src/custom/custom_modeling_mistral.py:537

Methodforward

( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = N

src/custom/custom_modeling_mistral.py:750

Methodforward

r""" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing

src/custom/custom_modeling_mistral.py:1009

Methodget_decoder

(self)

src/custom/custom_modeling_llama.py:1172

Methodget_decoder

(self)

src/custom/custom_modeling_mistral.py:1004

Methodget_input_embeddings

(self)

src/custom/custom_modeling_llama.py:945

Methodget_input_embeddings

(self)

src/custom/custom_modeling_llama.py:1157

Methodget_input_embeddings

(self)

src/custom/custom_modeling_mistral.py:743

Methodget_input_embeddings

(self)

src/custom/custom_modeling_mistral.py:989

Methodget_output_embeddings

(self)

src/custom/custom_modeling_llama.py:1163

Methodget_output_embeddings

(self)

src/custom/custom_modeling_mistral.py:995

Functionmean_pooling

(token_embeddings, mask)

experiments.py:47

Methodprepare_inputs_for_generation

( self, input_ids, past_key_values=None, attention_mask=None, inputs_e

src/custom/custom_modeling_llama.py:1270

Methodprepare_inputs_for_generation

( self, input_ids, past_key_values=None, attention_mask=None, inputs_e

src/custom/custom_modeling_mistral.py:1098

Methodrerank

Rerank the documents based on the query using a sliding window strategy. Assume that input documents are sorted by their relevance to

src/in_context_reranker.py:148

Methodset_decoder

(self, decoder)

src/custom/custom_modeling_llama.py:1169

Methodset_decoder

(self, decoder)

src/custom/custom_modeling_mistral.py:1001

Methodset_input_embeddings

(self, value)

src/custom/custom_modeling_llama.py:948

Methodset_input_embeddings

(self, value)

src/custom/custom_modeling_llama.py:1160

Methodset_input_embeddings

(self, value)

src/custom/custom_modeling_mistral.py:746

Methodset_input_embeddings

(self, value)

src/custom/custom_modeling_mistral.py:992

Methodset_output_embeddings

(self, new_embeddings)

src/custom/custom_modeling_llama.py:1166

Methodset_output_embeddings

(self, new_embeddings)

src/custom/custom_modeling_mistral.py:998