pytorch-conversational-memory-reranker

Conversational Memory Transformer Reranker

This is a small educational cross-encoder that scores whether a candidate conversational memory is relevant to a query. It demonstrates the complete PyTorch and Hugging Face workflow; it is not a production model.

Model details

• Architecture: BertForSequenceClassification with one output logit
• Base model: google/bert_uncased_L-2_H-128_A-2
• Size: approximately 4.4 million parameters in float32
• Language: English, uncased
• Input: one query and one candidate encoded together
• Maximum length: 128 WordPiece tokens, including special tokens
• Output: one raw relevance logit; larger values mean more predicted relevance

The inference example also applies sigmoid. Because the dataset is extremely small, that value is not a calibrated probability or reliable confidence score.

Intended uses

• Learning how transformer cross-encoders work
• Local experiments with query-candidate reranking
• Demonstrating training, evaluation, packaging, and model sharing
• Comparing a transformer with a simple bag-of-words baseline

Prohibited and out-of-scope uses

Do not use this model for production retrieval, safety-critical decisions, medical, legal, financial, employment, surveillance, identity, or access-control decisions. Do not use it to infer sensitive personal attributes. It has not been evaluated for languages other than English.

Training data

The dataset contains 50 entirely synthetic query-candidate examples organized into 10 query groups. Each group contains relevant, easy-negative, and hard-negative candidates. No real conversations, private data, patient data, company data, or user identifiers were used.

Query groups were split 60/20/20 into training, validation, and test data. The transformer was fine-tuned for four epochs with AdamW, a learning rate of 5e-5, and binary cross-entropy with logits. Epoch 2 was selected using validation MRR and nDCG.

Evaluation

The transformer did not outperform either baseline. These results cover only two synthetic test-query groups and do not establish general model quality.

Development-machine inference measured roughly 1.4 ms per pair and about 16.8 MiB for parameters and evaluated tensors. Timing varies by hardware and workload.

Minimal inference

python -m pip install -r requirements.txt
python inference.py . "What affects the user's sleep?" "Coffee keeps the user awake."

The script loads local package files and prints logit and probability fields.

Limitations and biases

• The fine-tuning dataset is far too small and synthetic for production use.
• The test set contains only two query groups.
• The model can rank irrelevant candidates above relevant ones.
• Sigmoid outputs are not calibrated probabilities.
• Inputs longer than 128 tokens are truncated and may lose key context.
• The model inherits biases and limitations from its pretrained base model.
• It has not been tested for robustness, fairness, adversarial inputs, or domain transfer.

Privacy and safety

The fine-tuning data is synthetic and contains no secrets or real personal data. The package contains model weights, public tokenizer files, configuration, documentation, and inference code only. Users must not supply data they are not authorized to process and should apply their own privacy controls.

Hardware and software

• Python 3.10 or newer
• PyTorch 2.13.0
• Transformers 5.16.1
• CPU inference is supported; a GPU is optional
• Approximately 17 MiB for float32 model parameters, plus framework overhead

License

This fine-tuned model is released under the Apache License 2.0, matching the license declared by the base model.