azza1625/target-identification

This model has been pushed to the Hub using the 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.

Citation

If you use this model in your research, please cite it as follows:

@article{charalampous2026compact,
  title={Compact Models for Structured Argument and Stance Analysis: A Framing-Aware, Retrieval-Augmented Pipeline},
  author={Charalampous, Antonis and Djouvas, Constantinos},
  journal={Machine Learning and Knowledge Extraction},
  volume={8},
  number={7},
  pages={203},
  year={2026},
  publisher={MDPI}
}

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:

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']}")