et_prediction_2 / README.md
skboy's picture
Update README.md
5785e77 verified
|
Raw
History Blame Contribute Delete
3.01 kB
# RoBERTa-based Eye-Tracking (ET) Feature Generator
## Overview
This repository contains the weights, tokenizer, and architecture for a custom regression model based on `roberta-base`. It is designed to predict 5 distinct eye-tracking (ET) features directly from text inputs. This model was trained to serve as the ET generator component required to replicate and extend the GazeReward framework.
## Reference
This model replicates the ET generator mentioned in the following work:
> "Through ablation studies we test our framework with different integration methods, LLMs, and ET generator models..."
> (Lopez-Cardona et al., "SEEING EYE TO AI: HUMAN ALIGNMENT VIA GAZE-BASED RESPONSE REWARDS FOR LARGE LANGUAGE MODELS")
## Model Architecture
- **Base Model:** `roberta-base`
- **Custom Head:** A linear layer that outputs 5 continuous ET features.
- **Implementation:** The exact architecture is defined in the accompanying `model.py` file.
## Training Data
The model was fine-tuned using eye-tracking data from:
- ZuCo 2.0 Dataset (CC BY-NC 4.0)
- Provo Corpus
## How to Use & Test
To use this model, download the weights (`.safetensors`) and the custom architecture script (`model.py`). You must use the `safetensors` library to load the weights. The tokenizer is included in this repository and can be loaded directly from the Hub.
Below is a complete script to download the model, load the weights and tokenizer, and run a quick inference test to verify everything works correctly. The code includes the required masking for special tokens.
```python
# File: test_inference.py
# Downloads the custom ET generator model and tokenizer from the Hugging Face Hub, loads them using safetensors, and runs a test inference.
import torch
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from model import RobertaRegressionModel
def run_quick_test(repo_id="skboy/et_prediction_2", filename="et_predictor2_seed123.safetensors"):
# Downloads weights, initializes the custom model, applies the tokenizer from the same repo, and prints the output tensor.
weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
model = RobertaRegressionModel()
state_dict = load_file(weights_path, device="cpu")
model.load_state_dict(state_dict)
model.eval()
tokenizer = AutoTokenizer.from_pretrained(repo_id)
sample_text = "This is a test sentence for eye-tracking feature generation."
inputs = tokenizer(sample_text, return_tensors="pt")
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
predict_mask = attention_mask.clone()
predict_mask[0, 0] = 0
predict_mask[0, -1] = 0
with torch.no_grad():
output = model(input_ids=input_ids, attention_mask=attention_mask, predict_mask=predict_mask)
print(f"Output shape: {output.shape}")
print(f"Output tensor:\n{output}")
if __name__ == "__main__":
run_quick_test()