Upload folder using huggingface_hub

README.md

@@ -0,0 +1,75 @@
+# AirRep-Flan
+
+AirRep is an attribution-friendly embedding model designed for computing training data influence on test examples.
+
+## Model Description
+
+This model is based on BERT architecture (gte-small config) with an additional projection layer. It's trained to produce embeddings that can be used for:
+- Text encoding
+- Computing similarity scores between test and training examples
+- Identifying influential training examples for test predictions
+
+## Model Details
+
+- **Base Architecture**: BERT (thenlper/gte-small config)
+- **Hidden Size**: 384
+- **Number of Layers**: 12
+- **Attention Heads**: 12
+- **Max Sequence Length**: 512
+- **Vocabulary Size**: 30522
+
+## Usage
+
+```python
+from airrep import AirRep
+
+# Load model
+model = AirRep.from_pretrained("sunweiwei/AirRep-Flan-Small")
+
+# Encode texts
+texts = ["Question: What is AI?\nAnswer: Artificial Intelligence..."]
+embeddings = model.encode(texts, batch_size=128, show_progress_bar=True)
+
+# Compute similarity scores
+test_embed = model.encode(test_texts)
+train_embed = model.encode(train_texts)
+scores = model.similarity(test_embed, train_embed, softmax=True)
+```
+
+## Installation
+
+```bash
+pip install airrep
+```
+
+Or install from source:
+
+```bash
+git clone https://github.com/sunnweiwei/AirRep
+cd AirRep
+pip install -e .
+```
+
+## Training Data
+
+This model was trained on the FLAN dataset with data influence optimization.
+
+## Evaluation
+
+- **Flan LDS Spearman Correlation**: 0.21
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@article{airrep2024,
+  title={AirRep: Attribution-friendly Representation Learning},
+  author={Sun, Weiwei},
+  year={2024}
+}
+```
+
+## License
+
+This model is released under the Apache 2.0 License.

config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "AirRepModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "airrep",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.57.1",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

model.safetensors

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

modeling_airrep.py

@@ -0,0 +1,95 @@
+"""AirRep model implementation."""
+
+from typing import Optional
+import torch
+import torch.nn as nn
+from transformers import BertModel, BertConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+
+def mean_pooling(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    """Apply mean pooling to hidden states."""
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+
+class AirRepConfig(BertConfig):
+    """Configuration class for AirRep model."""
+
+    model_type = "airrep"
+
+    def __init__(
+        self,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+
+class AirRepModel(PreTrainedModel):
+    """
+    AirRep model with BERT encoder and projection layer.
+
+    This is a standalone model, not a wrapper.
+    """
+
+    config_class = AirRepConfig
+    base_model_prefix = "airrep"
+
+    def __init__(self, config: AirRepConfig):
+        super().__init__(config)
+        self.config = config
+
+        # BERT encoder
+        self.bert = BertModel(config, add_pooling_layer=False)
+
+        # Projection layer
+        self.projector = nn.Linear(
+            config.hidden_size,
+            config.hidden_size,
+            dtype=torch.bfloat16
+        )
+
+        # Initialize weights
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            input_ids: Input token IDs
+            attention_mask: Attention mask
+            token_type_ids: Token type IDs
+
+        Returns:
+            Pooled and projected embeddings (batch_size, hidden_size)
+        """
+        # Get BERT outputs
+        outputs = self.bert(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        # Mean pooling
+        last_hidden_state = outputs.last_hidden_state
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        pooled = mean_pooling(last_hidden_state, attention_mask)
+
+        # Project
+        projected = self.projector(pooled)
+
+        return projected
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        """Save model and config."""
+        super().save_pretrained(save_directory, **kwargs)

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,65 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "max_length": 128,
+  "model_max_length": 1000000000000000019884624838656,
+  "never_split": null,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}