{"id":21874728,"url":"https://github.com/longxingtan/open-retrievals","last_synced_at":"2025-10-07T04:45:22.088Z","repository":{"id":226923884,"uuid":"769290116","full_name":"LongxingTan/open-retrievals","owner":"LongxingTan","description":"All-in-One: Text Embedding, Retrieval, Reranking and RAG in Transformers","archived":false,"fork":false,"pushed_at":"2025-05-23T04:16:45.000Z","size":1452,"stargazers_count":59,"open_issues_count":4,"forks_count":13,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-06-24T06:18:41.170Z","etag":null,"topics":["advanced-rag","contrastive-learning","embeddings","finetuning","information-retrieval","llm","llm-embeddings","llm-rerankers","nlp","rag","rag-rerank","rag-retrieval","retrieval","sentence-embedding","sentence-similarity","triplet-loss"],"latest_commit_sha":null,"homepage":"https://open-retrievals.readthedocs.io/en/master/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/LongxingTan.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2024-03-08T18:31:26.000Z","updated_at":"2025-06-07T17:56:35.000Z","dependencies_parsed_at":"2024-08-13T15:39:41.739Z","dependency_job_id":"5ee044ff-cc20-4f67-854b-ce48b313b369","html_url":"https://github.com/LongxingTan/open-retrievals","commit_stats":null,"previous_names":["longxingtan/open-retrievals"],"tags_count":8,"template":false,"template_full_name":null,"purl":"pkg:github/LongxingTan/open-retrievals","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LongxingTan%2Fopen-retrievals","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LongxingTan%2Fopen-retrievals/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LongxingTan%2Fopen-retrievals/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LongxingTan%2Fopen-retrievals/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/LongxingTan","download_url":"https://codeload.github.com/LongxingTan/open-retrievals/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LongxingTan%2Fopen-retrievals/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":261618263,"owners_count":23185108,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["advanced-rag","contrastive-learning","embeddings","finetuning","information-retrieval","llm","llm-embeddings","llm-rerankers","nlp","rag","rag-rerank","rag-retrieval","retrieval","sentence-embedding","sentence-similarity","triplet-loss"],"created_at":"2024-11-28T07:13:11.754Z","updated_at":"2025-10-07T04:45:17.070Z","avatar_url":"https://github.com/LongxingTan.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"[license-image]: https://img.shields.io/badge/License-Apache%202.0-blue.svg\n[license-url]: https://opensource.org/licenses/Apache-2.0\n[pypi-image]: https://badge.fury.io/py/open-retrievals.svg\n[pypi-url]: https://pypi.org/project/open-retrievals\n[pepy-image]: https://pepy.tech/badge/retrievals/month\n[pepy-url]: https://pepy.tech/project/retrievals\n[build-image]: https://github.com/LongxingTan/open-retrievals/actions/workflows/test.yml/badge.svg?branch=master\n[build-url]: https://github.com/LongxingTan/open-retrievals/actions/workflows/test.yml?query=branch%3Amaster\n[lint-image]: https://github.com/LongxingTan/open-retrievals/actions/workflows/lint.yml/badge.svg?branch=master\n[lint-url]: https://github.com/LongxingTan/open-retrievals/actions/workflows/lint.yml?query=branch%3Amaster\n[docs-image]: https://readthedocs.org/projects/open-retrievals/badge/?version=latest\n[docs-url]: https://open-retrievals.readthedocs.io/en/master/\n[coverage-image]: https://codecov.io/gh/longxingtan/open-retrievals/branch/master/graph/badge.svg\n[coverage-url]: https://codecov.io/github/longxingtan/open-retrievals?branch=master\n[contributing-image]: https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat\n[contributing-url]: https://github.com/longxingtan/open-retrievals/blob/master/CONTRIBUTING.md\n\n\u003ch1 align=\"center\"\u003e\n\u003cimg src=\"./docs/source/_static/logo.svg\" width=\"420\" align=center/\u003e\n\u003c/h1\u003e\n\n\u003cdiv align=\"center\"\u003e\n\n  [![LICENSE][license-image]][license-url]\n  [![PyPI Version][pypi-image]][pypi-url]\n  [![Build Status][build-image]][build-url]\n  [![Lint Status][lint-image]][lint-url]\n  [![Docs Status][docs-image]][docs-url]\n  [![Code Coverage][coverage-image]][coverage-url]\n  [![Contributing][contributing-image]][contributing-url]\n\n  **[Documentation](https://open-retrievals.readthedocs.io/en/master/)** | **[中文](https://github.com/LongxingTan/open-retrievals/blob/master/README_zh-CN.md)** | **[日本語](https://github.com/LongxingTan/open-retrievals/blob/master/README_ja-JP.md)**\n\n\u003c/div\u003e\n\n![structure](./docs/source/_static/structure.png)\n\n**Open-retrievals** unify text embedding, retrieval, reranking and RAG. It's easy, flexible and scalable to fine-tune the model.\n- Embedding fine-tuned through point-wise, pairwise, listwise, contrastive learning and LLM.\n- Reranking fine-tuned with Cross-Encoder, ColBERT and LLM.\n- Easily build enhanced modular RAG, integrated with Transformers, Langchain and LlamaIndex.\n\n| Experiment                    | Model                  | Original | Finetuned | Demo                                                                                                                                                                |\n|-------------------------------|------------------------|----------|-----------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| **embed** pairwise finetune   | bge-base-zh-v1.5       | 0.657    | **0.703** | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/17KXe2lnNRID-HiVvMtzQnONiO74oGs91?usp=sharing) |\n| **embed** LLM finetune (LoRA) | e5-mistral-7b-instruct | 0.651    | **0.699** | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1jj1kBQWFcuQ3a7P9ttnl1hgX7H8WA_Za?usp=sharing) |\n| **rerank** cross encoder      | bge-reranker-base      | 0.666    | **0.706** | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1QvbUkZtG56SXomGYidwI4RQzwODQrWNm?usp=sharing) |\n| **rerank** colbert            | bge-m3                 | 0.657    | **0.695** | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1QVtqhQ080ZMltXoJyODMmvEQYI6oo5kO?usp=sharing) |\n| **rerank** LLM (LoRA)         | bge-reranker-v2-gemma  | 0.637    | **0.706** | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1fzq1iV7-f8hNKFnjMmpVhVxadqPb9IXk?usp=sharing) |\n\n* The eval metrics is MAP in 10% [t2-reranking data](https://huggingface.co/datasets/C-MTEB/T2Reranking).\n* Read [more examples](./examples)\n\n\n## Installation\n\n**With pip**\n```shell\npip install transformers\npip install open-retrievals\n```\n\n\n## Quick-start\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1-WBMisdWLeHUKlzJ2DrREXY_kSV8vjP3?usp=sharing)\n\n\u003cdetails\u003e\u003csummary\u003e Embedding from pretrained weights \u003c/summary\u003e\n\n```python\nfrom retrievals import AutoModelForEmbedding\n\nsentences = [\n    'query: how much protein should a female eat',\n    'query: summit define',\n    \"passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\",\n    \"passage: Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.\"\n]\nmodel_name_or_path = 'intfloat/e5-base-v2'\nmodel = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method=\"mean\")\nembeddings = model.encode(sentences, normalize_embeddings=True)\nscores = (embeddings[:2] @ embeddings[2:].T) * 100\nprint(scores.tolist())\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\u003csummary\u003e Faiss retrieval search \u003c/summary\u003e\n\n```python\nfrom retrievals import AutoModelForEmbedding, AutoModelForRetrieval\n\nsentences = ['A dog is chasing car.', 'A man is playing a guitar.']\nmodel_name_or_path = \"sentence-transformers/all-MiniLM-L6-v2\"\nindex_path = './database/faiss/faiss.index'\nmodel = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method='mean')\nmodel.build_index(sentences, index_path=index_path)\n\nquery_embed = model.encode(\"He plays guitar.\")\nmatcher = AutoModelForRetrieval()\ndists, indices = matcher.search(query_embed, index_path=index_path)\nprint(indices)\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\u003csummary\u003e Rerank using pretrained weights \u003c/summary\u003e\n\n```python\nfrom retrievals import AutoModelForRanking\n\nmodel_name_or_path: str = \"BAAI/bge-reranker-base\"\nrerank_model = AutoModelForRanking.from_pretrained(model_name_or_path)\nscores_list = rerank_model.compute_score([\"In 1974, I won the championship in Southeast Asia in my first kickboxing match\", \"In 1982, I defeated the heavy hitter Ryu Long.\"])\nprint(scores_list)\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\u003csummary\u003e RAG with LangChain integration \u003c/summary\u003e\n\n```shell\npip install langchain\npip install chromadb langchain-chroma\n```\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1fJC-8er-a4NRkdJkwWr4On7lGt9rAO4P?usp=sharing)\n\n```python\nfrom retrievals.tools.langchain import LangchainEmbedding, LangchainReranker, LangchainLLM\nfrom retrievals import AutoModelForRanking\nfrom langchain.retrievers import ContextualCompressionRetriever\nfrom langchain.prompts.prompt import PromptTemplate\nfrom langchain.chains import RetrievalQA\nfrom langchain_chroma import Chroma as Vectorstore\n\npersist_directory = './database/faiss.index'\nembed_model_name_or_path = \"sentence-transformers/all-MiniLM-L6-v2\"\nrerank_model_name_or_path = \"BAAI/bge-reranker-base\"\nllm_model_name_or_path = \"microsoft/Phi-3-mini-128k-instruct\"\n\nembeddings = LangchainEmbedding(model_name=embed_model_name_or_path)\nvectordb = Vectorstore(\n    collection_name=\"example_collection\",\n    persist_directory=persist_directory,\n    embedding_function=embeddings,\n)\nretrieval_args = {\"search_type\" :\"similarity\", \"score_threshold\": 0.15, \"k\": 10}\nretriever = vectordb.as_retriever(**retrieval_args)\n\nranker = AutoModelForRanking.from_pretrained(rerank_model_name_or_path)\nreranker = LangchainReranker(model=ranker, top_n=3)\ncompression_retriever = ContextualCompressionRetriever(\n    base_compressor=reranker, base_retriever=retriever\n)\n\nllm = LangchainLLM(model_name_or_path=llm_model_name_or_path)\n\nRESPONSE_TEMPLATE = \"\"\"[INST]\n\u003c\u003e\nYou are a helpful AI assistant. Use the following pieces of context to answer the user's question.\u003c\u003e\nAnything between the following `context` html blocks is retrieved from a knowledge base.\n\n    {context}\n\nREMEMBER:\n- If you don't know the answer, just say that you don't know, don't try to make up an answer.\n- Let's take a deep breath and think step-by-step.\n\nQuestion: {question}[/INST]\nHelpful Answer:\n\"\"\"\n\nPROMPT = PromptTemplate(template=RESPONSE_TEMPLATE, input_variables=[\"context\", \"question\"])\n\nqa_chain = RetrievalQA.from_chain_type(\n    llm,\n    chain_type='stuff',\n    retriever=compression_retriever,\n    chain_type_kwargs={\n        \"verbose\": True,\n        \"prompt\": PROMPT,\n    }\n)\n\nuser_query = 'Introduce this'\nresponse = qa_chain({\"query\": user_query})\nprint(response)\n```\n\u003c/details\u003e\n\n\n## Fine-tuning\n\n\u003cdetails\u003e\u003csummary\u003e Fine-tune embedding \u003c/summary\u003e\n\n```python\nimport torch.nn as nn\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments\nfrom retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel\nfrom retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss\n\nmodel_name_or_path: str = \"sentence-transformers/paraphrase-multilingual-mpnet-base-v2\"\nbatch_size: int = 32\nepochs: int = 3\n\ntrain_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']\ntokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)\nmodel = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method=\"mean\")\ntrain_model = PairwiseModel(model)\n\noptimizer = AdamW(train_model.parameters(), lr=5e-5)\nnum_train_steps = int(len(train_dataset) / batch_size * epochs)\nscheduler = get_linear_schedule_with_warmup(\n    optimizer, num_warmup_steps=0.05 * num_train_steps, num_training_steps=num_train_steps\n)\n\ntraining_arguments = TrainingArguments(\n    output_dir='./checkpoints',\n    num_train_epochs=epochs,\n    per_device_train_batch_size=batch_size,\n    remove_unused_columns=False,\n    logging_steps=100,\n    report_to=\"none\",\n)\ntrainer = RetrievalTrainer(\n    model=train_model,\n    args=training_arguments,\n    train_dataset=train_dataset,\n    data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),\n    loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)),\n)\ntrainer.optimizer = optimizer\ntrainer.scheduler = scheduler\ntrainer.train()\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\u003csummary\u003e Fine-tune LLM embedding \u003c/summary\u003e\n\n```python\nimport torch.nn as nn\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments\nfrom retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel\nfrom retrievals.losses import InfoNCE, SimCSE, TripletLoss\n\ndef add_instructions(example):\n    example['query'] = query_instruction.format(example['query'])\n    example['positive'] = document_instruction.format(example['positive'])\n    return example\n\nmodel_name_or_path: str = \"Qwen/Qwen2-1.5B-Instruct\"\nbatch_size: int = 8\nepochs: int = 3\nquery_instruction = \"Retrieve relevant passages that answer the query\\nQuery: {}\"\ndocument_instruction = \"Document: {}\"\n\ntrain_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']\ntrain_dataset = train_dataset.rename_columns({'sentence1': 'query', 'sentence2': 'positive'})\ntrain_dataset = train_dataset.map(add_instructions)\ntokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)\nmodel = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method=\"last\", use_lora=True)\ntrain_model = PairwiseModel(model, loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)))\noptimizer = AdamW(train_model.parameters(), lr=5e-5)\nnum_train_steps = int(len(train_dataset) / batch_size * epochs)\nscheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0.05 * num_train_steps, num_training_steps=num_train_steps)\n\ntraining_arguments = TrainingArguments(\n    output_dir='./checkpoints',\n    num_train_epochs=epochs,\n    per_device_train_batch_size=batch_size,\n    remove_unused_columns=False,\n    logging_steps=100,\n    report_to=\"none\",\n)\ntrainer = RetrievalTrainer(\n    model=train_model,\n    args=training_arguments,\n    train_dataset=train_dataset,\n    data_collator=RetrievalCollator(tokenizer, keys=['query', 'positive'], max_lengths=[64, 128]),\n)\ntrainer.optimizer = optimizer\ntrainer.scheduler = scheduler\ntrainer.train()\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\u003csummary\u003e Fine-tune cross-encoder reranking \u003c/summary\u003e\n\n```python\nfrom transformers import AutoTokenizer, TrainingArguments, get_cosine_schedule_with_warmup, AdamW\nfrom retrievals import RerankCollator, AutoModelForRanking, RerankTrainer, RerankTrainDataset\n\nmodel_name_or_path: str = \"BAAI/bge-reranker-base\"\nmax_length: int = 128\nlearning_rate: float = 3e-5\nbatch_size: int = 4\nepochs: int = 3\noutput_dir: str = \"./checkpoints\"\n\ntrain_dataset = RerankTrainDataset(\"C-MTEB/T2Reranking\", positive_key=\"positive\", negative_key=\"negative\", dataset_split='dev')\ntokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)\nmodel = AutoModelForRanking.from_pretrained(model_name_or_path)\noptimizer = AdamW(model.parameters(), lr=learning_rate)\nnum_train_steps = int(len(train_dataset) / batch_size * epochs)\nscheduler = get_cosine_schedule_with_warmup(\n    optimizer,\n    num_warmup_steps=0.05 * num_train_steps,\n    num_training_steps=num_train_steps,\n)\n\ntraining_args = TrainingArguments(\n    learning_rate=learning_rate,\n    per_device_train_batch_size=batch_size,\n    num_train_epochs=epochs,\n    output_dir=output_dir,\n    remove_unused_columns=False,\n    logging_steps=100,\n    report_to=\"none\",\n)\ntrainer = RerankTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=train_dataset,\n    data_collator=RerankCollator(tokenizer, max_length=max_length),\n)\ntrainer.optimizer = optimizer\ntrainer.scheduler = scheduler\ntrainer.train()\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\u003csummary\u003e Fine-tune ColBERT reranking \u003c/summary\u003e\n\n```python\nimport os\nimport transformers\nfrom transformers import (\n    AdamW,\n    AutoTokenizer,\n    TrainingArguments,\n    get_cosine_schedule_with_warmup,\n)\n\nfrom retrievals import ColBERT, ColBertCollator, RerankTrainer, RetrievalTrainDataset\nfrom retrievals.losses import ColbertLoss\n\ntransformers.logging.set_verbosity_error()\nos.environ[\"WANDB_DISABLED\"] = \"true\"\n\nmodel_name_or_path: str = \"BAAI/bge-m3\"\nlearning_rate: float = 5e-6\nbatch_size: int = 32\nepochs: int = 3\ncolbert_dim: int = 1024\noutput_dir: str = './checkpoints'\n\ntrain_dataset = RetrievalTrainDataset('C-MTEB/T2Reranking', positive_key='positive', negative_key='negative', dataset_split='dev')\ntokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)\ndata_collator = ColBertCollator(\n    tokenizer,\n    query_max_length=128,\n    document_max_length=256,\n    positive_key='positive',\n    negative_key='negative',\n)\nmodel = ColBERT.from_pretrained(\n    model_name_or_path,\n    colbert_dim=colbert_dim,\n    loss_fn=ColbertLoss(use_inbatch_negative=False),\n)\n\noptimizer = AdamW(model.parameters(), lr=learning_rate)\nnum_train_steps = int(len(train_dataset) / batch_size * epochs)\nscheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=0.05 * num_train_steps, num_training_steps=num_train_steps)\n\ntraining_args = TrainingArguments(\n    learning_rate=learning_rate,\n    per_device_train_batch_size=batch_size,\n    num_train_epochs=epochs,\n    output_dir=output_dir,\n    remove_unused_columns=False,\n    logging_steps=100,\n    report_to=\"none\",\n)\ntrainer = RerankTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=train_dataset,\n    data_collator=data_collator,\n)\ntrainer.optimizer = optimizer\ntrainer.scheduler = scheduler\ntrainer.train()\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\u003csummary\u003e Fine-tune LLM reranking \u003c/summary\u003e\n\n```python\nfrom transformers import (\n    AdamW,\n    AutoTokenizer,\n    TrainingArguments,\n    get_cosine_schedule_with_warmup,\n)\n\nfrom retrievals import (\n    LLMRanker,\n    LLMRerankCollator,\n    RerankTrainer,\n    RetrievalTrainDataset,\n)\nfrom retrievals.losses import TokenLoss\n\nmodel_name_or_path: str = \"Qwen/Qwen2-1.5B-Instruct\"\nmax_length: int = 512\nlearning_rate: float = 3e-5\nbatch_size: int = 8\nepochs: int = 3\ntask_prompt: str = (\n    \"\"\"Given a query A and a passage B, determine whether the passage contains an answer to the query\"\"\"\n    \"\"\"by providing a prediction of either 'Yes' or 'No'.\"\"\"\n)\n\ntokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)\ntrain_dataset = RetrievalTrainDataset(\n    data_name_or_path='C-MTEB/T2Reranking',\n    positive_key='positive',\n    negative_key='negative',\n    query_instruction='A: {}',\n    document_instruction='B: {}',\n    dataset_split='dev',\n)\ndata_collator = LLMRerankCollator(tokenizer=tokenizer, max_length=max_length, prompt=task_prompt, add_target_token='Yes')\ntoken_index = tokenizer('Yes', add_special_tokens=False)['input_ids'][-1]\nmodel = LLMRanker.from_pretrained(\n    model_name_or_path,\n    causal_lm=True,\n    use_fp16=True,\n    loss_fn=TokenLoss(token_index=token_index),\n    use_lora=True,\n)\n\noptimizer = AdamW(model.parameters(), lr=learning_rate)\nnum_train_steps = int(len(train_dataset) / batch_size * epochs)\nscheduler = get_cosine_schedule_with_warmup(\n    optimizer,\n    num_warmup_steps=0.05 * num_train_steps,\n    num_training_steps=num_train_steps,\n)\n\ntraining_args = TrainingArguments(\n    learning_rate=learning_rate,\n    per_device_train_batch_size=batch_size,\n    num_train_epochs=epochs,\n    output_dir=\"./checkpoints\",\n    remove_unused_columns=False,\n    report_to=\"none\",\n)\ntrainer = RerankTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=train_dataset,\n    data_collator=data_collator,\n)\ntrainer.optimizer = optimizer\ntrainer.scheduler = scheduler\ntrainer.train()\n```\n\u003c/details\u003e\n\n\n## RAG module\ncoming soon\n\n\n## Reference \u0026 Acknowledge\n- [UKPLab/sentence-transformers](https://github.com/UKPLab/sentence-transformers)\n- [luyug/Dense](https://github.com/luyug/Dense)\n- [FlagOpen/FlagEmbedding](https://github.com/FlagOpen/FlagEmbedding)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flongxingtan%2Fopen-retrievals","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flongxingtan%2Fopen-retrievals","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flongxingtan%2Fopen-retrievals/lists"}