pip uninstall -y torch torchvision torchaudio
pip install --no-cache-dir torch torchvision torchaudio \
  --index-url https://download.pytorch.org/whl/cu118

python

# Load model directly
from transformers import AutoModel

device_type = 'cuda'
model = AutoModel.from_pretrained("HiteshSaai/arxiv_research-cat-classifier-adapter", dtype="auto")

input_sequence = {'content': 'TITLE: Beyond the Finite Variant Property: Extending Symbolic Diffie-Hellman Group Models (Extended Version) \n Abstract: Diffie-Hellman groups are commonly used in cryptographic protocols. While most state-of-the-art, symbolic protocol verifiers support them to some degr \n\n Task: Classify this paper into one research category.\nAnswer with only the category name.',
  'role': 'user'}

prompt = tokenizer.apply_chat_template(input_sequence, tokenize=False, add_generation_prompt=True)

tokenized_input = tokenizer(prompt, return_tensors='pt', add_special_tokens=False).to(device_type)

out = model.generate(**tokenized_input, max_new_tokens=5)

response = tokenizer.batch_decode(out, skip_special_tokens=False)

classified_output = response[0][len(prompt):]

Training procedure

This model was trained with SFT.

Framework versions

TRL: 0.15.2
Transformers: 4.49.0
Pytorch: 2.6.0
Datasets: 3.3.2
Tokenizers: 0.21.4

Citations

Cite TRL as:

bibtex

@misc{vonwerra2022trl,
	title        = {{TRL: Transformer Reinforcement Learning}},
	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallouédec},
	year         = 2020,
	journal      = {GitHub repository},
	publisher    = {GitHub},
	howpublished = {\url{https://github.com/huggingface/trl}}
}