🇺🇦 Keyword Extraction via Bi-Encoder Retrieval (Ukrainian)

📌 Overview

This repository contains a sentence-transformers bi-encoder fine-tuned for extracting keywords from Ukrainian scientific and business documents.

Unlike a generative LLM, this model does not write keywords — it retrieves them. The document is encoded into the same vector space as a pre-built index of candidate keywords, and the closest keywords by cosine similarity are returned. This makes the model much smaller and faster than a generative LoRA-tuned LLM, while reaching very high retrieval accuracy on the evaluation set (see Evaluation).

ℹ️ This replaces the previous generative approach (LoRA adapter on meta-llama/Llama-3.1-8B-Instruct) with a dedicated retrieval model.

🧠 Base Model

• Base model: intfloat/multilingual-e5-base
• Architecture: XLM-RoBERTa encoder + mean pooling + L2 normalization
• Fine-tuning method: Full fine-tuning with MultipleNegativesRankingLoss
• Library: Sentence Transformers
• Output dimensionality: 768
• Max sequence length: 512 tokens
• Similarity function: Cosine similarity

Full Model Architecture

SentenceTransformer(
  (0): Transformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'XLMRobertaModel'})
  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_mean_tokens': True})
  (2): Normalize()
)

📊 Dataset & Training

• Training pairs: 78,648 (document, keyword) samples built from a corpus of Ukrainian scientific/business documents (see dataset_full.jsonl used at training time — not included in this repo, only the resulting keyword index is).
• Query/passage format follows the E5 convention:
  • Documents are prefixed with query:
  • Keywords are prefixed with passage:
• Loss: MultipleNegativesRankingLoss (in-batch negatives, scale=20.0)
• Epochs: 3, batch size 96 (grad. accumulation 2), learning rate 2e-5, bf16

Training Logs

📈 Evaluation

Evaluated with InformationRetrievalEvaluator on held-out (document, keyword) pairs:

🚀 Usage

1. Raw embeddings (Sentence Transformers)

pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("ipsnan/keywordsextraction")

sentences = [
    "query: Штучний інтелект активно використовується у промисловості для автоматизації процесів...",
    "passage: штучний інтелект",
    "passage: автоматизація процесів",
]
embeddings = model.encode(sentences)
print(embeddings.shape)  # [3, 768]

similarities = model.similarity(embeddings, embeddings)
print(similarities)

2. Full keyword extraction pipeline (PDF → keywords)

This repository also ships keyword_index.pkl — embeddings of every unique keyword seen during training — and infer_keywords.py, the retrieval script used to go from a raw PDF document straight to a ranked keyword list.

pip install -U sentence-transformers pymupdf numpy

python infer_keywords.py /path/to/file.pdf
python infer_keywords.py /path/to/file.pdf --top-k 10 --threshold 0.5

Under the hood the script:

1. Extracts and cleans text from the PDF (1fullsourcedata.py).
2. Encodes the first ~3000 characters of the document with this model, prefixed with query: .
3. Looks up the nearest keyword vectors in keyword_index.pkl by cosine similarity.
4. Returns the top-K keywords above --threshold.

To rebuild the keyword index from scratch you need the original dataset_full.jsonl training corpus and can run python infer_keywords.py <pdf> --rebuild-index.

🎯 Task

Keyword extraction (Ukrainian NLP), optimized for:

• scientific texts
• academic content
• business documents

⚠️ Limitations

• Works best with Ukrainian-language texts.
• Can only return keywords that already exist in keyword_index.pkl (it is a retrieval model, not a generative one) — rebuild the index on your own keyword set to extend coverage.
• Performance depends on how representative the first ~3000 characters of the document are (title + abstract).

📦 Repository Contents

• model.safetensors, config.json, config_sentence_transformers.json, modules.json, sentence_bert_config.json, tokenizer.json, tokenizer_config.json, 1_Pooling/, 2_Normalize/ — the Sentence Transformers model
• keyword_index.pkl — pre-computed keyword embeddings used for retrieval
• infer_keywords.py — end-to-end PDF → keywords inference script
• 1fullsourcedata.py — PDF text extraction/cleaning helper used by infer_keywords.py

📜 License

• Model weights: MIT
• Base model: subject to the intfloat/multilingual-e5-base license (MIT)

🙌 Acknowledgments

• intfloat/multilingual-e5-base
• Sentence Transformers / Hugging Face ecosystem

Download model

Download them in the Files & versions tab.