---
tags:
- target-identification
- argumentation
- contrastive-learning
license: mit
language:
- en
base_model:
- answerdotai/ModernBERT-base
pipeline_tag: text-classification
---

This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:

---
## Model Description
This is a dual-encoder retrieval model built on top of `answerdotai/ModernBERT-base`. The model is designed to perform target identification by finding the most relevant `theses` along with their associated data for a given `claim`

You can modify the `top_k`, `num_args` & `top_level_only` variables to adjust the output of the model.

## How to use
You can use this model for inference by loading it with the `transformers` library. The following code demonstrates how to make a prediction:

```python
import torch
import torch.nn as nn

from transformers import AutoModel, AutoTokenizer
from huggingface_hub import hf_hub_download, PyTorchModelHubMixin

import pickle
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

class DualEncoderThesisModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self) -> None:
        super(DualEncoderThesisModel, self).__init__()
        self.encoder = AutoModel.from_pretrained("answerdotai/ModernBERT-base")

    def forward(self, input_ids_a, attention_mask_a, input_ids_b, attention_mask_b):
        # Encode arguments
        output_a = self.encoder(input_ids=input_ids_a, attention_mask=attention_mask_a).last_hidden_state
        emb_a = output_a[:, 0]

        # Encode theses
        output_b = self.encoder(input_ids=input_ids_b, attention_mask=attention_mask_b).last_hidden_state
        emb_b = output_b[:, 0]

        return emb_a, emb_b

model_name = "azza1625/target-identification"
tokenizer = AutoTokenizer.from_pretrained(model_name)

model = DualEncoderThesisModel.from_pretrained(model_name)
model.eval()

device = "cpu"

embeddings_path = hf_hub_download(
    repo_id="azza1625/target-identification",
    filename="retrieval_data_random_negatives_10_train_data.pkl"
)

with open(embeddings_path, "rb") as f:
    embeddings_metadata = pickle.load(f)

@torch.no_grad()
def retrieve_theses(claim, top_k=3, num_args=5, top_level_only=True, device="cpu"):
    stored_embeddings = embeddings_metadata["embeddings"]
    metadata = embeddings_metadata["metadata"]
    
    enc = tokenizer(claim, return_tensors='pt', truncation=True, padding='max_length', max_length=1024).to(device)
    query_embedding = model.encoder(**enc).last_hidden_state[:, 0].cpu().numpy()

    sims = cosine_similarity(query_embedding, stored_embeddings)[0]
    top_indices = np.argsort(sims)[::-1][:top_k]

    results = []
    for idx in top_indices:
        arguments = metadata[idx]['arguments']
        if top_level_only:
            arguments = [arg for arg in arguments if arg['target_type'] == 'thesis']

        results.append({
            "thesis": metadata[idx]["thesis"],
            "debate_title": metadata[idx]["debate_title"],
            "arguments": arguments[:num_args]
        })

    return results

claim = "A fetus or embryo is not a person; therefore, abortion should not be considered murder."

theses = retrieve_theses(claim)

for thesis in theses:
    print(f"{thesis['thesis']} | {thesis['debate_title']}")