├── .env
├── .gitignore
├── .pre-commit-config.yaml
├── CITATION.cff
├── LICENSE
├── Makefile
├── NOTICE.txt
├── README.md
├── __init__.py
├── inc_rel
    ├── __init__.py
    ├── args.py
    ├── datasets.py
    ├── download.py
    ├── eval.py
    ├── few_shot_trainer.py
    ├── first_stage.py
    ├── generate_few_shot.py
    ├── index.py
    ├── knn_eval.py
    ├── knn_index.py
    ├── knn_similarities.py
    ├── meta
    │   ├── __init__.py
    │   ├── adapter_meta.py
    │   ├── bias_meta.py
    │   ├── classification_head_meta.py
    │   ├── learner.py
    │   ├── meta_dataset.py
    │   ├── module_meta.py
    │   └── utils.py
    ├── pre_train.py
    ├── pre_train_trainer.py
    ├── query_fine_tune.py
    ├── rank_fusion.py
    ├── reranking_evaluator.py
    ├── settings.py
    ├── utils.py
    └── zero_shot.py
├── requirements.dev.txt
└── requirements.txt


/.env:
--------------------------------------------------------------------------------
 1 | IR_DATASETS_HOME=~/.ir_datasets
 2 | 
 3 | # General Settings
 4 | INC_REL_BASE_PATH=/Volumes/SD/data/inc-rel
 5 | INC_REL_RAW_DIR=raw
 6 | INC_REL_DATA_DIR=data
 7 | INC_REL_ES_DIR=es
 8 | 
 9 | # TREC-COVID
10 | INC_REL_TREC_COVID_RAW_PATH=${INC_REL_BASE_PATH}/${INC_REL_RAW_DIR}/trec-covid
11 | INC_REL_TREC_COVID_DATA_PATH=${INC_REL_BASE_PATH}/${INC_REL_DATA_DIR}/trec-covid
12 | INC_REL_TREC_COVID_CORPUS_URL=https://ai2-semanticscholar-cord-19.s3-us-west-2.amazonaws.com/historical_releases/cord-19_2020-07-16.tar.gz
13 | INC_REL_TREC_COVID_TOPICS_URL=https://ir.nist.gov/trec-covid/data/topics-rnd5.xml
14 | INC_REL_TREC_COVID_QRELS_URL=https://ir.nist.gov/trec-covid/data/qrels-covid_d5_j0.5-5.txt
15 | 
16 | # ROBUST04
17 | INC_REL_TREC_ROBUST_RAW_PATH=${INC_REL_BASE_PATH}/${INC_REL_RAW_DIR}/robust04
18 | INC_REL_TREC_ROBUST_DATA_PATH=${INC_REL_BASE_PATH}/${INC_REL_DATA_DIR}/robust04
19 | INC_REL_TREC_ROBUST_USERNAME=<NONE>
20 | INC_REL_TREC_ROBUST_PASSWORD=<NONE>
21 | INC_REL_TREC_ROBUST_CORPUS_DISK4_URL=https://ir.nist.gov/cd45/TREC-Disk-4.tar.gz
22 | INC_REL_TREC_ROBUST_CORPUS_DISK5_URL=https://ir.nist.gov/cd45/TREC-Disk-5.tar.gz
23 | 
24 | # TREC-NEWS
25 | INC_REL_TREC_NEWS_RAW_PATH=${INC_REL_BASE_PATH}/${INC_REL_RAW_DIR}/trec-news
26 | INC_REL_TREC_NEWS_DATA_PATH=${INC_REL_BASE_PATH}/${INC_REL_DATA_DIR}/trec-news
27 | INC_REL_TREC_NEWS_USERNAME=<NONE>
28 | INC_REL_TREC_NEWS_PASSWORD=<NONE>
29 | INC_REL_TREC_NEWS_CORPUS_URL=https://ir.nist.gov/wapo/WashingtonPost.v2.tar.gz
30 | INC_REL_TREC_NEWS_TOPICS_URL=https://trec.nist.gov/data/news/2019/newsir19-background-linking-topics.xml
31 | INC_REL_TREC_NEWS_QRELS_URL=https://trec.nist.gov/data/news/2019/newsir19-qrels-background.txt
32 | 
33 | # WEBIS-TOUCHE-2020
34 | INC_REL_TOUCHE_IR_DATASETS_NAME=beir/webis-touche2020
35 | INC_REL_TOUCHE_DATA_PATH=${INC_REL_BASE_PATH}/${INC_REL_DATA_DIR}/touche
36 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | .venv
2 | .DS_Store
3 | __pycache__
4 | *.pyc
5 | .python-version
6 | elasticsearch-*
7 | 


--------------------------------------------------------------------------------
/.pre-commit-config.yaml:
--------------------------------------------------------------------------------
 1 | repos:
 2 | -   repo: https://github.com/psf/black
 3 |     rev: 22.6.0
 4 |     hooks:
 5 |     - id: black
 6 |       args: [--check]
 7 |       language_version: python3.10
 8 | -   repo: https://github.com/pycqa/isort
 9 |     rev: 5.10.1
10 |     hooks:
11 |       - id: isort
12 |         name: isort (python)
13 |         args: ["--profile", "black"]
14 | 


--------------------------------------------------------------------------------
/CITATION.cff:
--------------------------------------------------------------------------------
 1 | # This CITATION.cff file was generated with cffinit.
 2 | # Visit https://bit.ly/cffinit to generate yours today!
 3 | 
 4 | cff-version: 1.2.0
 5 | title: >-
 6 |   Incorporating Relevance Feedback for
 7 |   Information-Seeking Retrieval using Few-Shot
 8 |   Document Re-Ranking
 9 | message: >-
10 |   If you use this software, please cite it using the
11 |   metadata from this file.
12 | type: software
13 | authors:
14 |   - given-names: Tim
15 |     family-names: Baumgärtner
16 |     email: baumgaertner.t@gmail.com
17 |     affiliation: TU Darmstadt
18 |   - given-names: 'Leonardo F. R. '
19 |     family-names: Ribeiro
20 |     affiliation: Amazon Alexa AI
21 |   - given-names: Nils
22 |     family-names: Reimers
23 |     affiliation: cohere.ai
24 |   - given-names: Iryna
25 |     family-names: Gurevych
26 |     affiliation: TU Darmstadt
27 | abstract: >-
28 |   Pairing a lexical retriever with a neural
29 |   re-ranking model has set state-of-the-art
30 |   performance on large-scale information retrieval
31 |   datasets. This pipeline covers scenarios like
32 |   question answering or navigational queries,
33 |   however, for information-seeking scenarios, users
34 |   often provide information on whether a document is
35 |   relevant to their query in form of clicks or
36 |   explicit feedback. Therefore, in this work, we
37 |   explore how relevance feedback can be directly
38 |   integrated into neural re-ranking models by
39 |   adopting few-shot and parameter-efficient learning
40 |   techniques. Specifically, we introduce a kNN
41 |   approach that re-ranks documents based on their
42 |   similarity with the query and the documents the
43 |   user considers relevant. Further, we explore
44 |   Cross-Encoder models that we pre-train using
45 |   meta-learning and subsequently fine-tune for each
46 |   query, training only on the feedback documents. To
47 |   evaluate our different integration strategies, we
48 |   transform four existing information retrieval
49 |   datasets into the relevance feedback scenario.
50 |   Extensive experiments demonstrate that integrating
51 |   relevance feedback directly in neural re-ranking
52 |   models improves their performance, and fusing
53 |   lexical ranking with our best performing neural
54 |   re-ranker outperforms all other methods by 5.2
55 |   nDCG@20.
56 | repository-code: https://github.com/UKPLab/incorporating-relevance
57 | preferred-citation:
58 |   type: proceedings
59 |   authors:
60 |   - given-names: Tim
61 |     family-names: Baumgärtner
62 |   - given-names: 'Leonardo F. R. '
63 |     family-names: Ribeiro
64 |   - given-names: Nils
65 |     family-names: Reimers
66 |   - given-names: Iryna
67 |     family-names: Gurevych
68 |   title: >-
69 |     Incorporating Relevance Feedback for
70 |     Information-Seeking Retrieval using Few-Shot
71 |     Document Re-Ranking
72 |   year: 2022
73 |   url: https://arxiv.org/abs/2210.10695
74 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
  1 |                                  Apache License
  2 |                            Version 2.0, January 2004
  3 |                         http://www.apache.org/licenses/
  4 | 
  5 |    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
  6 | 
  7 |    1. Definitions.
  8 | 
  9 |       "License" shall mean the terms and conditions for use, reproduction,
 10 |       and distribution as defined by Sections 1 through 9 of this document.
 11 | 
 12 |       "Licensor" shall mean the copyright owner or entity authorized by
 13 |       the copyright owner that is granting the License.
 14 | 
 15 |       "Legal Entity" shall mean the union of the acting entity and all
 16 |       other entities that control, are controlled by, or are under common
 17 |       control with that entity. For the purposes of this definition,
 18 |       "control" means (i) the power, direct or indirect, to cause the
 19 |       direction or management of such entity, whether by contract or
 20 |       otherwise, or (ii) ownership of fifty percent (50%) or more of the
 21 |       outstanding shares, or (iii) beneficial ownership of such entity.
 22 | 
 23 |       "You" (or "Your") shall mean an individual or Legal Entity
 24 |       exercising permissions granted by this License.
 25 | 
 26 |       "Source" form shall mean the preferred form for making modifications,
 27 |       including but not limited to software source code, documentation
 28 |       source, and configuration files.
 29 | 
 30 |       "Object" form shall mean any form resulting from mechanical
 31 |       transformation or translation of a Source form, including but
 32 |       not limited to compiled object code, generated documentation,
 33 |       and conversions to other media types.
 34 | 
 35 |       "Work" shall mean the work of authorship, whether in Source or
 36 |       Object form, made available under the License, as indicated by a
 37 |       copyright notice that is included in or attached to the work
 38 |       (an example is provided in the Appendix below).
 39 | 
 40 |       "Derivative Works" shall mean any work, whether in Source or Object
 41 |       form, that is based on (or derived from) the Work and for which the
 42 |       editorial revisions, annotations, elaborations, or other modifications
 43 |       represent, as a whole, an original work of authorship. For the purposes
 44 |       of this License, Derivative Works shall not include works that remain
 45 |       separable from, or merely link (or bind by name) to the interfaces of,
 46 |       the Work and Derivative Works thereof.
 47 | 
 48 |       "Contribution" shall mean any work of authorship, including
 49 |       the original version of the Work and any modifications or additions
 50 |       to that Work or Derivative Works thereof, that is intentionally
 51 |       submitted to Licensor for inclusion in the Work by the copyright owner
 52 |       or by an individual or Legal Entity authorized to submit on behalf of
 53 |       the copyright owner. For the purposes of this definition, "submitted"
 54 |       means any form of electronic, verbal, or written communication sent
 55 |       to the Licensor or its representatives, including but not limited to
 56 |       communication on electronic mailing lists, source code control systems,
 57 |       and issue tracking systems that are managed by, or on behalf of, the
 58 |       Licensor for the purpose of discussing and improving the Work, but
 59 |       excluding communication that is conspicuously marked or otherwise
 60 |       designated in writing by the copyright owner as "Not a Contribution."
 61 | 
 62 |       "Contributor" shall mean Licensor and any individual or Legal Entity
 63 |       on behalf of whom a Contribution has been received by Licensor and
 64 |       subsequently incorporated within the Work.
 65 | 
 66 |    2. Grant of Copyright License. Subject to the terms and conditions of
 67 |       this License, each Contributor hereby grants to You a perpetual,
 68 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 69 |       copyright license to reproduce, prepare Derivative Works of,
 70 |       publicly display, publicly perform, sublicense, and distribute the
 71 |       Work and such Derivative Works in Source or Object form.
 72 | 
 73 |    3. Grant of Patent License. Subject to the terms and conditions of
 74 |       this License, each Contributor hereby grants to You a perpetual,
 75 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 76 |       (except as stated in this section) patent license to make, have made,
 77 |       use, offer to sell, sell, import, and otherwise transfer the Work,
 78 |       where such license applies only to those patent claims licensable
 79 |       by such Contributor that are necessarily infringed by their
 80 |       Contribution(s) alone or by combination of their Contribution(s)
 81 |       with the Work to which such Contribution(s) was submitted. If You
 82 |       institute patent litigation against any entity (including a
 83 |       cross-claim or counterclaim in a lawsuit) alleging that the Work
 84 |       or a Contribution incorporated within the Work constitutes direct
 85 |       or contributory patent infringement, then any patent licenses
 86 |       granted to You under this License for that Work shall terminate
 87 |       as of the date such litigation is filed.
 88 | 
 89 |    4. Redistribution. You may reproduce and distribute copies of the
 90 |       Work or Derivative Works thereof in any medium, with or without
 91 |       modifications, and in Source or Object form, provided that You
 92 |       meet the following conditions:
 93 | 
 94 |       (a) You must give any other recipients of the Work or
 95 |           Derivative Works a copy of this License; and
 96 | 
 97 |       (b) You must cause any modified files to carry prominent notices
 98 |           stating that You changed the files; and
 99 | 
100 |       (c) You must retain, in the Source form of any Derivative Works
101 |           that You distribute, all copyright, patent, trademark, and
102 |           attribution notices from the Source form of the Work,
103 |           excluding those notices that do not pertain to any part of
104 |           the Derivative Works; and
105 | 
106 |       (d) If the Work includes a "NOTICE" text file as part of its
107 |           distribution, then any Derivative Works that You distribute must
108 |           include a readable copy of the attribution notices contained
109 |           within such NOTICE file, excluding those notices that do not
110 |           pertain to any part of the Derivative Works, in at least one
111 |           of the following places: within a NOTICE text file distributed
112 |           as part of the Derivative Works; within the Source form or
113 |           documentation, if provided along with the Derivative Works; or,
114 |           within a display generated by the Derivative Works, if and
115 |           wherever such third-party notices normally appear. The contents
116 |           of the NOTICE file are for informational purposes only and
117 |           do not modify the License. You may add Your own attribution
118 |           notices within Derivative Works that You distribute, alongside
119 |           or as an addendum to the NOTICE text from the Work, provided
120 |           that such additional attribution notices cannot be construed
121 |           as modifying the License.
122 | 
123 |       You may add Your own copyright statement to Your modifications and
124 |       may provide additional or different license terms and conditions
125 |       for use, reproduction, or distribution of Your modifications, or
126 |       for any such Derivative Works as a whole, provided Your use,
127 |       reproduction, and distribution of the Work otherwise complies with
128 |       the conditions stated in this License.
129 | 
130 |    5. Submission of Contributions. Unless You explicitly state otherwise,
131 |       any Contribution intentionally submitted for inclusion in the Work
132 |       by You to the Licensor shall be under the terms and conditions of
133 |       this License, without any additional terms or conditions.
134 |       Notwithstanding the above, nothing herein shall supersede or modify
135 |       the terms of any separate license agreement you may have executed
136 |       with Licensor regarding such Contributions.
137 | 
138 |    6. Trademarks. This License does not grant permission to use the trade
139 |       names, trademarks, service marks, or product names of the Licensor,
140 |       except as required for reasonable and customary use in describing the
141 |       origin of the Work and reproducing the content of the NOTICE file.
142 | 
143 |    7. Disclaimer of Warranty. Unless required by applicable law or
144 |       agreed to in writing, Licensor provides the Work (and each
145 |       Contributor provides its Contributions) on an "AS IS" BASIS,
146 |       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 |       implied, including, without limitation, any warranties or conditions
148 |       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 |       PARTICULAR PURPOSE. You are solely responsible for determining the
150 |       appropriateness of using or redistributing the Work and assume any
151 |       risks associated with Your exercise of permissions under this License.
152 | 
153 |    8. Limitation of Liability. In no event and under no legal theory,
154 |       whether in tort (including negligence), contract, or otherwise,
155 |       unless required by applicable law (such as deliberate and grossly
156 |       negligent acts) or agreed to in writing, shall any Contributor be
157 |       liable to You for damages, including any direct, indirect, special,
158 |       incidental, or consequential damages of any character arising as a
159 |       result of this License or out of the use or inability to use the
160 |       Work (including but not limited to damages for loss of goodwill,
161 |       work stoppage, computer failure or malfunction, or any and all
162 |       other commercial damages or losses), even if such Contributor
163 |       has been advised of the possibility of such damages.
164 | 
165 |    9. Accepting Warranty or Additional Liability. While redistributing
166 |       the Work or Derivative Works thereof, You may choose to offer,
167 |       and charge a fee for, acceptance of support, warranty, indemnity,
168 |       or other liability obligations and/or rights consistent with this
169 |       License. However, in accepting such obligations, You may act only
170 |       on Your own behalf and on Your sole responsibility, not on behalf
171 |       of any other Contributor, and only if You agree to indemnify,
172 |       defend, and hold each Contributor harmless for any liability
173 |       incurred by, or claims asserted against, such Contributor by reason
174 |       of your accepting any such warranty or additional liability.
175 | 
176 |    END OF TERMS AND CONDITIONS
177 | 
178 |    APPENDIX: How to apply the Apache License to your work.
179 | 
180 |       To apply the Apache License to your work, attach the following
181 |       boilerplate notice, with the fields enclosed by brackets "[]"
182 |       replaced with your own identifying information. (Don't include
183 |       the brackets!)  The text should be enclosed in the appropriate
184 |       comment syntax for the file format. We also recommend that a
185 |       file or class name and description of purpose be included on the
186 |       same "printed page" as the copyright notice for easier
187 |       identification within third-party archives.
188 | 
189 |    Copyright [yyyy] [name of copyright owner]
190 | 
191 |    Licensed under the Apache License, Version 2.0 (the "License");
192 |    you may not use this file except in compliance with the License.
193 |    You may obtain a copy of the License at
194 | 
195 |        http://www.apache.org/licenses/LICENSE-2.0
196 | 
197 |    Unless required by applicable law or agreed to in writing, software
198 |    distributed under the License is distributed on an "AS IS" BASIS,
199 |    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 |    See the License for the specific language governing permissions and
201 |    limitations under the License.
202 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
  1 | include .env
  2 | 
  3 | .PHONY: install install-dev format download index first-stage zero-shot knn query-fine-tune meta-query-fine-tune rank-fusion help
  4 | .DEFAULT_GOAL := help
  5 | DOCKER_AVAILABLE := $(shell dockerdd -v 2>/dev/null)
  6 | ES_DIR := ./elasticsearch-7.11.2
  7 | 
  8 | ## 		Install dependencies. Note: Activate virtualenv before running this command.
  9 | install:
 10 | 	pip install --upgrade pip
 11 | 	pip install -r requirements.txt
 12 | 
 13 | ## 		Install development dependencies. Note: Activate virtualenv before running this command.
 14 | install-dev: 
 15 | 	pip install --upgrade pip
 16 | 	pip install -r requirements.dev.txt
 17 | 	pre-commit install
 18 | 
 19 | ## 		Format code base.
 20 | format: 
 21 | 	isort --profile black .
 22 | 	black .
 23 | 	
 24 | ## 		Download datasets according to the configuration in .env.
 25 | download: 
 26 | 	python inc_rel/download.py
 27 | 
 28 | es-up: 
 29 | ifdef DOCKER_AVAILABLE
 30 | 	@echo Runnig Elasticsearch in docker
 31 | 	docker run -d \
 32 | 		-v $(INC_REL_ES_DATA_DIR):/usr/share/elasticsearch/data \
 33 | 		-p 9200:9200 \
 34 | 		-e "discovery.type=single-node" \
 35 | 		-e "indices.query.bool.max_clause_count=16384" \
 36 | 		--name inc-rel-es \
 37 | 		elasticsearch:7.11.2
 38 | else
 39 | 	@echo Running Elasticsearch on host
 40 | 	$(ES_DIR)/bin/elasticsearch -Ediscovery.type=single-node -Eindices.query.bool.max_clause_count=16438 -d -p $(PWD)/es.pid
 41 | endif
 42 | 
 43 | 	bash -c "until curl -s -o /dev/null http://localhost:9200; do echo 'Waiting for Elasticsearch'; sleep 3; done"
 44 | 
 45 | es-down:
 46 | ifdef DOCKER_AVAILABLE
 47 | 	docker rm -s inc-rel-es
 48 | else
 49 | 	kill `cat es.pid`
 50 | endif
 51 | 
 52 | generate-first-stage:
 53 | 	IR_DATASETS_HOME=$(IR_DATASETS_HOME) python inc_rel/first_stage.py --dataset $(dataset) || $(MAKE) es-down
 54 | 	
 55 | first-stage: dataset:=$(dataset)
 56 | ## 		Get first stage retrieval results using BM25.
 57 | first-stage: es-up generate-first-stage es-down
 58 | 
 59 | generate-few-shot:
 60 | 	IR_DATASETS_HOME=$(IR_DATASETS_HOME) python inc_rel/generate_few_shot.py --dataset $(dataset) || $(MAKE) es-down
 61 | index: dataset:=$(dataset)
 62 | ## 		Create an elasticsearch index; perform first and second stage retrieval and generate the few-shot dataset.
 63 | index: es-up generate-few-shot es-down
 64 | 
 65 | 
 66 | ## 		Evaluate zero-shot re-ranking.
 67 | zero-shot: 
 68 | 	python inc_rel/zero_shot.py $(args)
 69 | 
 70 | knn-index: 
 71 | 	python inc_rel/knn_index.py $(args)
 72 | knn-similarities: 
 73 | 	python inc_rel/knn_similarities.py $(args)
 74 | knn-eval: 
 75 | 	python inc_rel/knn_eval.py $(args)
 76 | knn: args:=$(args)
 77 | ## 			Evaluate knn re-ranking.
 78 | knn: knn-index knn-similarities knn-eval
 79 | 	@echo args=$(args)
 80 | 
 81 | ## 		Fine-tune the re-ranker per query on the few-shot examples.
 82 | query-ft:
 83 | 	python inc_rel/query_fine_tune.py $(args)
 84 | 
 85 | ## 	Pre-train the re-ranker using supervised or meta learning, then Fine-tune per query on the few-shot examples.
 86 | pre-train-query-ft:
 87 | 	python inc_rel/pre_train.py $(args)
 88 | 
 89 | ## 		Fuse two or more rankings together.
 90 | rank-fusion:
 91 | 	python inc_rel/rank_fusion.py $(args)
 92 | 
 93 | # COLORS
 94 | GREEN  := $(shell tput -Txterm setaf 2)
 95 | YELLOW := $(shell tput -Txterm setaf 3)
 96 | RESET  := $(shell tput -Txterm sgr0)
 97 | 
 98 | ## 			Show this help.	
 99 | help:
100 | 	@echo ''
101 | 	@echo 'Usage:'
102 | 	@echo '  ${YELLOW}make${RESET} ${GREEN}<target>${RESET}'
103 | 	@echo ''
104 | 	@echo 'Targets:'
105 | 	@awk '/^[a-zA-Z\-\_0-9]+:/ { \
106 | 		helpMessage = match(lastLine, /^## (.*)/); \
107 | 		if (helpMessage) { \
108 | 			helpCommand = substr($$1, 0, index($$1, ":")-1); \
109 | 			helpMessage = substr(lastLine, RSTART + 3, RLENGTH); \
110 | 			printf "  ${YELLOW}%-$(TARGET_MAX_CHAR_NUM)s${RESET} ${GREEN}%s${RESET}\n", helpCommand, helpMessage; \
111 | 		} \
112 | 	} \
113 | 	{ lastLine = $$0 }' $(MAKEFILE_LIST)
114 | 


--------------------------------------------------------------------------------
/NOTICE.txt:
--------------------------------------------------------------------------------
1 | -------------------------------------------------------------------------------
2 | Copyright 2022
3 | Ubiquitous Knowledge Processing (UKP) Lab
4 | Technische Universität Darmstadt
5 | -------------------------------------------------------------------------------
6 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # Incorporating Relevance Feedback for Information-Seeking Retrieval using Few-Shot Document Re-Ranking
  2 | <p align="center">
  3 | <a href="https://aclanthology.org/2022.emnlp-main.614/">
  4 |     <img alt="Paper Badge" src="https://img.shields.io/badge/Paper-b31b1b.svg">
  5 | </a>
  6 | <a href="https://public.ukp.informatik.tu-darmstadt.de/baumgaertner/emnlp-2022/inc-rel-slides-emnlp-2022.pdf">
  7 |     <img alt="Slides Badge" src="https://img.shields.io/badge/Slides-b31b1b.svg">
  8 | </a>
  9 | <a href="https://public.ukp.informatik.tu-darmstadt.de/baumgaertner/emnlp-2022/inc-rel-poster-emnlp-2022.pdf">
 10 |     <img alt="Slides Badge" src="https://img.shields.io/badge/Poster-b31b1b.svg">
 11 | </a>
 12 | </p>
 13 | 
 14 | This repository contains the code for our EMNLP 2022 paper [_Incorporating Relevance Feedback for Information-Seeking Retrieval using Few-Shot Document Re-Ranking_](https://aclanthology.org/2022.emnlp-main.614).
 15 | 
 16 | ## Overview
 17 | <p align="center">
 18 | <img src="https://user-images.githubusercontent.com/11020443/214001736-5aa508ac-56cb-4c89-8bf9-ba879693f868.svg">
 19 | </p>
 20 | 
 21 | > Pairing a lexical retriever with a neural re-ranking model has set state-of-the-art performance on large-scale information retrieval datasets. This pipeline covers scenarios like question answering or navigational queries, however, for information-seeking scenarios, users often provide information on whether a document is relevant to their query in form of clicks or explicit feedback. Therefore, in this work, we explore how relevance feedback can be directly integrated into neural re-ranking models by adopting few-shot and parameter-efficient learning techniques. Specifically, we introduce a kNN approach that re-ranks documents based on their similarity with the query and the documents the user considers relevant. Further, we explore Cross-Encoder models that we pre-train using meta-learning and subsequently fine-tune for each query, training only on the feedback documents. To evaluate our different integration strategies, we transform four existing information retrieval datasets into the relevance feedback scenario. Extensive experiments demonstrate that integrating relevance feedback directly in neural re-ranking models improves their performance, and fusing lexical ranking with our best performing neural re-ranker outperforms all other methods by 5.2 nDCG@20.
 22 | 
 23 | ## Requirements
 24 | - python 3.10+
 25 | - if docker is availbale, elasticsearch will be run in a container, else the bare metal version will be started.
 26 | ## Setup
 27 | ### Virtual Environment and Dependencies
 28 | 1. Setup your python virtual environment, for example:
 29 |     ```shell
 30 |     python -m venv .venv
 31 |     ```
 32 | 2. Install dependencies
 33 |     ```shell
 34 |     make install
 35 |     ```
 36 |     If you want to contribute to the repository, please also install the dev dependencies:
 37 |     ```shell
 38 |     make install-dev
 39 |     ```
 40 | ### Datasets and Preprocessing
 41 | The paths where to store raw, preprocessed and experiments data can be modified in the `.env` file. Once these are set, run the following commands:
 42 | 1. Download Datasets
 43 |     ```shell
 44 |     make download
 45 |     ```
 46 | 2. Index corpora in elasticsearch.
 47 |     ```
 48 |     make index
 49 |     ```
 50 | 
 51 | 
 52 | ## Experiments
 53 | The following commands run the experiments using the default parameters. For the full set of available arguments, see [inc_rel/args.py](inc_rel/args.py).
 54 | ### 2nd Stage Retrieval and Query Expansion
 55 | Creating the index the last step in the [Setup](#setup) also runs the experiments over the second stage retrieval. By default, the query is expanded with `[4, 8, 16, 32, 64]` terms using ElasticSearchs [MoreLikeThis](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-mlt-query.html) query.
 56 | 
 57 | ### kNN
 58 | ```
 59 | make knn args="--dataset <dataset>"
 60 | ```
 61 | 
 62 | ### Zero-Shot
 63 | ```
 64 | make knn args="--dataset <dataset> --num_samples <num_samples>"
 65 | ```
 66 | 
 67 | ### Query Fine-Tuning
 68 | ```
 69 | make args="--dataset <dataset> --num_samples <num_samples>"
 70 | ```
 71 | 
 72 | ### Meta-Learning + Query Fine-Tuning
 73 | ```
 74 | make args="--dataset <dataset> --num_samples <num_samples>"
 75 | ```
 76 | 
 77 | ### Rank-Fusion
 78 | ```
 79 | make args="--dataset <dataset> --num_samples <num_samples> --result_files <path to first result file> <path to second result file> 
 80 | ```
 81 | 
 82 | ## Contact
 83 | This project is maintained by [Tim Baumgärtner](https://github.com/timbmg).
 84 | - [Mail](mailto:baumgaertner.t@gmail.com) 
 85 | - [Twitter](https://twitter.com/timbmg) 
 86 | - [UKP Lab](http://www.ukp.tu-darmstadt.de/)
 87 | - [TU Darmstadt](http://www.tu-darmstadt.de/)
 88 | 
 89 | ## Citation
 90 | If you find this work useful, please considering citing the following [paper](https://aclanthology.org/2022.emnlp-main.614):
 91 | ```bibtex
 92 | @inproceedings{baumgartner-etal-2022-incorporating,
 93 |     title = "Incorporating Relevance Feedback for Information-Seeking Retrieval using Few-Shot Document Re-Ranking",
 94 |     author = {Baumg{\"a}rtner, Tim  and
 95 |       Ribeiro, Leonardo F. R.  and
 96 |       Reimers, Nils  and
 97 |       Gurevych, Iryna},
 98 |     booktitle = "Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing",
 99 |     month = dec,
100 |     year = "2022",
101 |     address = "Abu Dhabi, United Arab Emirates",
102 |     publisher = "Association for Computational Linguistics",
103 |     url = "https://aclanthology.org/2022.emnlp-main.614",
104 |     pages = "8988--9005",
105 | }
106 | ```
107 | 
108 | ## Disclaimer
109 | This repository contains experimental software and is published for the sole purpose of giving additional background details on the respective publication. 
110 | 


--------------------------------------------------------------------------------
/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/UKPLab/incorporating-relevance/c7419c136f0b07bfe105095ebb49d2bb088c00b6/__init__.py


--------------------------------------------------------------------------------
/inc_rel/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/UKPLab/incorporating-relevance/c7419c136f0b07bfe105095ebb49d2bb088c00b6/inc_rel/__init__.py


--------------------------------------------------------------------------------
/inc_rel/args.py:
--------------------------------------------------------------------------------
  1 | import json
  2 | import os
  3 | from dataclasses import dataclass, field
  4 | from enum import Enum
  5 | from typing import Dict, List, Union
  6 | 
  7 | from sentence_transformers import CrossEncoder, SentenceTransformer
  8 | from settings import dataset_settings_cls
  9 | 
 10 | 
 11 | class ScoringFunction(str, Enum):
 12 |     cos = "cos"
 13 |     dot = "dot"
 14 | 
 15 | 
 16 | class FTParams(str, Enum):
 17 |     full = "full"
 18 |     bias = "bias"
 19 |     # adapter = "adapter"
 20 |     head = "head"
 21 | 
 22 | 
 23 | class PreTrainMethod(str, Enum):
 24 |     supervised = "supervised"
 25 |     meta = "meta"
 26 | 
 27 | 
 28 | @dataclass(kw_only=True)
 29 | class Experiment:
 30 |     exp_name: str = "exp"
 31 |     dataset: str = list(dataset_settings_cls.keys())[0]
 32 |     num_samples: int = 2
 33 |     seeds: List[int] = field(default_factory=lambda: [0, 1, 2])
 34 |     splits: List[str] = field(default_factory=lambda: ["train", "valid", "test"])
 35 |     bm25_size: int = 1000
 36 |     metric: str = "ndcg_cut_20"
 37 | 
 38 |     @property
 39 |     def dataset_settings(self):
 40 |         return dataset_settings_cls[self.dataset]()
 41 | 
 42 |     @property
 43 |     def data_path(self) -> str:
 44 |         return os.path.join(
 45 |             self.dataset_settings.data_path, str(self.dataset_settings.bm25_size)
 46 |         )
 47 | 
 48 |     @property
 49 |     def exp_path(self) -> str:
 50 |         _exp_path = os.path.join(
 51 |             self.dataset_settings.data_path, "experiments", self.exp_name
 52 |         )
 53 |         os.makedirs(_exp_path, exist_ok=True)
 54 |         return _exp_path
 55 | 
 56 |     @property
 57 |     def bm25_results(self) -> Dict:
 58 |         if not hasattr(self, "_bm25_results"):
 59 |             file = os.path.join(
 60 |                 self.data_path, f"k{self.num_samples}", "expansion_results_16.json"
 61 |             )
 62 |             with open(file) as fh:
 63 |                 self._bm25_results = json.load(fh)
 64 | 
 65 |         return self._bm25_results
 66 | 
 67 |     @property
 68 |     def bm25_docs(self) -> Dict:
 69 |         if not hasattr(self, "_bm25_docs"):
 70 |             file = os.path.join(
 71 |                 self.data_path, f"k{self.num_samples}", "expansion_docs_16.json"
 72 |             )
 73 |             with open(file) as fh:
 74 |                 self._bm25_docs = json.load(fh)
 75 | 
 76 |         return self._bm25_docs
 77 | 
 78 |     @property
 79 |     def qrels(self) -> Dict:
 80 |         if not hasattr(self, "_qrels"):
 81 |             file = os.path.join(self.data_path, "qrels.json")
 82 |             with open(file) as fh:
 83 |                 self._qrels = json.load(fh)
 84 | 
 85 |         return self._qrels
 86 | 
 87 |     @property
 88 |     def topics(self) -> Dict:
 89 |         if not hasattr(self, "_topics"):
 90 |             file = os.path.join(self.data_path, "topics.json")
 91 |             with open(file) as fh:
 92 |                 self._topics = json.load(fh)
 93 | 
 94 |         return self._topics
 95 | 
 96 |     @property
 97 |     def topic_ids_split_seed(self) -> Dict:
 98 |         if not hasattr(self, "_split_seed"):
 99 |             self._split_seed = {}
100 |             for seed in self.seeds:
101 |                 for split in self.splits:
102 |                     file = os.path.join(
103 |                         self.data_path,
104 |                         f"k{self.num_samples}",
105 |                         f"s{seed}",
106 |                         f"{split}.json",
107 |                     )
108 |                     with open(file) as fh:
109 |                         self._split_seed[split, seed] = json.load(fh)
110 | 
111 |         return self._split_seed
112 | 
113 | 
114 | @dataclass(kw_only=True)
115 | class ZeroShot(Experiment):
116 |     exp_name: str = "zero-shot"
117 |     model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"
118 |     model_ctx: Union[SentenceTransformer, None] = None
119 |     scoring_fn: ScoringFunction = "cos"
120 | 
121 |     @property
122 |     def model_class(self) -> str:
123 |         if self.model.startswith("cross-encoder"):
124 |             _model_class = CrossEncoder
125 |         else:
126 |             _model_class = SentenceTransformer
127 |         return _model_class
128 | 
129 | 
130 | @dataclass(kw_only=True)
131 | class KNN(Experiment):
132 |     num_samples: List[int] = field(default_factory=lambda: [2, 4, 8])
133 |     exp_name: str = "knn"
134 |     model: str = "sentence-transformers/all-MiniLM-L6-v2"
135 |     scoring_fn: ScoringFunction = "cos"
136 | 
137 | 
138 | @dataclass(kw_only=True)
139 | class FineTuneExperiment(Experiment):
140 |     exp_name: str = "query-ft"
141 |     model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"
142 |     ft_params: FTParams = "bias"
143 |     eval_batch_size: int = 32
144 |     epochs: int = 8
145 |     learning_rates: List[float] = field(default_factory=lambda: [2e-3, 2e-4, 2e-5])
146 | 
147 |     @property
148 |     def model_class(self) -> str:
149 |         if self.model.startswith("cross-encoder"):
150 |             _model_class = CrossEncoder
151 |         else:
152 |             _model_class = SentenceTransformer
153 |         return _model_class
154 | 
155 |     @property
156 |     def hparam_results_file(self) -> str:
157 |         return os.path.join(
158 |             self.exp_path,
159 |             f"k{self.num_samples}_s{{seed}}_valid_{self.ft_params}_hpsearch.json",
160 |         )
161 | 
162 | 
163 | @dataclass(kw_only=True)
164 | class PreTrain(FineTuneExperiment):
165 |     exp_name: str = "pt-query-ft"
166 |     pretrain_method: PreTrainMethod = "meta"
167 | 
168 | 
169 | @dataclass(kw_only=True)
170 | class RankFusion(Experiment):
171 |     exp_name: str = "rf"
172 |     result_files: List[str]
173 | 


--------------------------------------------------------------------------------
/inc_rel/datasets.py:
--------------------------------------------------------------------------------
  1 | import csv
  2 | import itertools
  3 | import json
  4 | import xml.etree.ElementTree as ET
  5 | from collections import defaultdict
  6 | from pathlib import Path
  7 | from typing import Callable, Dict, List, Union
  8 | 
  9 | import ir_datasets
 10 | import utils
 11 | from tqdm.auto import tqdm
 12 | 
 13 | 
 14 | class Dataset:
 15 |     def __init__(
 16 |         self,
 17 |         corpus_path: Path,
 18 |         topics_path: Path,
 19 |         qrels_path: Path,
 20 |         load_corpus: bool,
 21 |         remove_topics_with_few_annotations: bool = True,
 22 |         remove_topics_with_few_negatives: bool = True,
 23 |     ):
 24 |         if load_corpus:
 25 |             self.corpus = self.read_corpus(corpus_path)
 26 |         self.topics = self.read_topics(topics_path)
 27 |         self.qrels = self.read_qrels(qrels_path)
 28 |         if remove_topics_with_few_annotations:
 29 |             self.remove_topics_with_few_annotations(remove_topics_with_few_negatives)
 30 |         self.corpus_size = None
 31 | 
 32 |     def remove_topic(self, topic_id):
 33 |         self.qrels.pop(topic_id, None)
 34 |         self.topics.pop(topic_id, None)
 35 | 
 36 |     def read_corpus(self, corpus_path: Path) -> Union[Dict, Callable]:
 37 |         raise NotImplementedError()
 38 | 
 39 |     def read_topics(self, topics_path: Path) -> Dict[str, str]:
 40 |         raise NotImplementedError()
 41 | 
 42 |     def read_qrels(self, qrels_path: Path) -> Dict[str, str]:
 43 |         raise NotImplementedError()
 44 | 
 45 |     def remove_topics_with_few_annotations(
 46 |         self, remove_topics_with_few_negatives: bool
 47 |     ):
 48 | 
 49 |         topics_to_remove = []
 50 |         for topic_id, doc_id2relevance in self.qrels.items():
 51 |             num_relevant = len(
 52 |                 list(filter(lambda r: r >= 1, doc_id2relevance.values()))
 53 |             )
 54 |             if remove_topics_with_few_negatives:
 55 |                 num_non_relevant = len(
 56 |                     list(filter(lambda r: r <= 0, doc_id2relevance.values()))
 57 |                 )
 58 |             else:
 59 |                 num_non_relevant = self.min_annotations
 60 |             if self.min_annotations > min(num_relevant, num_non_relevant):
 61 |                 topics_to_remove.append((topic_id, num_relevant, num_non_relevant))
 62 | 
 63 |         for topic_id in self.topics.keys():
 64 |             if topic_id not in self.qrels:
 65 |                 topics_to_remove.append((topic_id, 0, 0))
 66 | 
 67 |         for topic_id, num_relevant, num_non_relevant in topics_to_remove:
 68 |             print(
 69 |                 f"Removing {topic_id=} due to too few annotations "
 70 |                 f"({num_relevant=}, {num_non_relevant=})."
 71 |             )
 72 |             self.remove_topic(topic_id)
 73 | 
 74 |         if not len(self.qrels) == len(self.topics):
 75 |             raise RuntimeError(
 76 |                 "Expected to have same amount of qrels and topics, "
 77 |                 f"but got {len(self.qrels)=} {len(self.topics)=}"
 78 |             )
 79 |         print(f"Remaining Topics={len(self.qrels)}")
 80 | 
 81 |     def remove_annotations_from_qrels(self, annotations: Dict[str, List[Dict]]):
 82 |         for topic_id, annotation in annotations.items():
 83 |             topic_id = annotation[0]["topic_id"]
 84 |             for a in annotation:
 85 |                 doc_id = a["doc_id"]
 86 |                 r = self.qrels[topic_id].pop(doc_id, None)
 87 |                 if r is None:
 88 |                     print(
 89 |                         f"Could not find document to remove from qrels. {topic_id=} {doc_id=} label={a['label']}"
 90 |                     )
 91 | 
 92 | 
 93 | class TrecCovid(Dataset):
 94 |     def __init__(self, *args, **kwargs):
 95 |         self.min_annotations = 32
 96 |         self.min_relevant_relevancy = 2
 97 |         self.corpus_size = 192509
 98 |         super().__init__(*args, **kwargs)
 99 | 
100 |     def read_corpus(self, corpus_path: Path) -> Dict[str, str]:
101 | 
102 |         fields = ["title", "abstract"]
103 | 
104 |         corpus = {}
105 |         with open(corpus_path, encoding="utf8") as fh:
106 |             reader = csv.DictReader(fh, delimiter=",", quoting=csv.QUOTE_MINIMAL)
107 |             for row in reader:
108 |                 cord_uid = row["cord_uid"]
109 |                 row["title"] = row["title"].strip()
110 |                 if row["title"] and row["title"][-1] not in "!?.":
111 |                     row["title"] += "."
112 | 
113 |                 text = " ".join([row[field] for field in fields]).strip()
114 |                 if cord_uid and text:
115 |                     corpus[cord_uid] = text
116 | 
117 |         return corpus
118 | 
119 |     def read_topics(self, topics_path: Path) -> Dict[str, str]:
120 | 
121 |         query_type = "question"
122 | 
123 |         queries = {}
124 |         root = ET.parse(topics_path).getroot()
125 |         for topic in root:
126 |             for query in topic:
127 |                 if query.tag == query_type:
128 |                     queries[topic.attrib["number"]] = query.text
129 | 
130 |         return queries
131 | 
132 |     def read_qrels(self, qrels_path: Path) -> Dict[str, Dict[str, int]]:
133 |         qrels = defaultdict(dict)
134 |         with open(qrels_path, "r") as fh:
135 |             for line in fh:
136 |                 query_id, _, doc_id, relevance = line.strip().strip("\n").split(" ")
137 |                 qrels[str(query_id)][doc_id] = int(relevance)
138 | 
139 |         return qrels
140 | 
141 | 
142 | class TrecNews(Dataset):
143 |     def __init__(self, *args, **kwargs):
144 |         self.min_annotations = 32
145 |         self.min_relevant_relevancy = 1
146 |         self.corpus_size = 595037
147 |         super().__init__(*args, **kwargs)
148 | 
149 |     def read_topics(self, topics_path: Path) -> Dict[str, str]:
150 | 
151 |         topics = {}
152 |         with open(topics_path) as fh:
153 |             it = itertools.chain("<root>", fh, "</root>")
154 |             root = ET.fromstringlist(it)
155 | 
156 |         docid2num = {}
157 |         for topic in root.findall("top"):
158 |             num = topic[0].text.split(":")[1].strip()
159 |             docid = topic[1].text
160 |             docid2num[docid] = num
161 |         for doc in tqdm(
162 |             self.corpus(doc_ids=list(docid2num.keys()), verbose=False),
163 |             desc="extracting titles",
164 |             total=len(docid2num),
165 |             ncols=100,
166 |         ):
167 |             num = docid2num[doc["id"]]
168 |             topics[num] = doc["title"]
169 | 
170 |         return topics
171 | 
172 |     def read_qrels(self, qrels_path: Path) -> Dict[str, Dict[str, int]]:
173 |         qrels = defaultdict(dict)
174 |         with open(qrels_path) as fh:
175 |             for line in fh.readlines():
176 |                 qid, _, doc_id, relevance = line.split(" ")
177 |                 qrels[str(qid)].update({doc_id: int(relevance)})
178 |         return qrels
179 | 
180 |     def read_corpus(self, corpus_path: Path) -> Callable:
181 |         def corpus_iter(doc_ids: List[str] = None, verbose=True):
182 |             if doc_ids is not None:
183 |                 doc_ids = list(set(doc_ids))
184 |             with open(corpus_path) as fh:
185 |                 for line in tqdm(
186 |                     fh, total=self.corpus_size, disable=not verbose, ncols=100
187 |                 ):
188 |                     doc = json.loads(line)
189 |                     if doc_ids is None or doc["id"] in doc_ids:
190 |                         full_text = []
191 |                         title = ""
192 |                         for content in doc.get("contents", []):
193 |                             if content is None:
194 |                                 continue
195 |                             if content["type"] == "title":
196 |                                 title = content["content"]
197 |                             elif (
198 |                                 content.get("type", None) == "sanitized_html"
199 |                                 and content.get("subtype", None) == "paragraph"
200 |                             ):
201 |                                 full_text.append(utils.strip_tags(content["content"]))
202 |                         full_text = " ".join(full_text)
203 | 
204 |                         yield {"id": doc["id"], "title": title, "text": full_text}
205 | 
206 |                         if doc_ids is not None:
207 |                             doc_ids.remove(doc["id"])
208 |                             if len(doc_ids) == 0:
209 |                                 break
210 | 
211 |         return corpus_iter
212 | 
213 | 
214 | class Robust04(Dataset):
215 |     def __init__(self, *args, **kwargs):
216 |         self.min_annotations = 32
217 |         self.min_relevant_relevancy = 1
218 |         self.dataset = ir_datasets.load("trec-robust04")
219 |         self.corpus_size = 528155
220 |         super().__init__(*args, **kwargs)
221 | 
222 |     def read_corpus(self, *args, **kwargs) -> Dict[str, str]:
223 |         corpus = {}
224 |         for doc in tqdm(
225 |             self.dataset.docs_iter(),
226 |             total=self.corpus_size,
227 |             desc="loading dataset",
228 |             ncols=100,
229 |         ):
230 |             corpus[doc.doc_id] = doc.text
231 |         return corpus
232 | 
233 |     def read_topics(self, *args, **kwargs) -> Dict[str, str]:
234 |         topics = {}
235 |         for topic in self.dataset.queries_iter():
236 |             topics[topic.query_id] = topic.description
237 |         return topics
238 | 
239 |     def read_qrels(self, *args, **kwargs) -> Dict[str, Dict[str, int]]:
240 |         qrels = defaultdict(dict)
241 |         for qrel in self.dataset.qrels_iter():
242 |             qrels[qrel.query_id][qrel.doc_id] = qrel.relevance
243 |         return qrels
244 | 
245 | 
246 | class Touche(Dataset):
247 |     def __init__(self, *args, **kwargs):
248 |         self.min_annotations = 32
249 |         self.min_relevant_relevancy = 3
250 |         self.dataset = ir_datasets.load("beir/webis-touche2020")
251 |         self.corpus_size = 382545
252 |         super().__init__(*args, **kwargs)
253 | 
254 |     def read_corpus(self, *args, **kwargs) -> Dict[str, str]:
255 |         corpus = {}
256 |         for doc in tqdm(
257 |             self.dataset.docs_iter(),
258 |             total=self.corpus_size,
259 |             desc="loading dataset",
260 |             ncols=100,
261 |         ):
262 |             corpus[doc.doc_id] = doc.text
263 |         return corpus
264 | 
265 |     def read_topics(self, *args, **kwargs) -> Dict[str, str]:
266 |         topics = {}
267 |         for topic in self.dataset.queries_iter():
268 |             topics[topic.query_id] = topic.text
269 |         return topics
270 | 
271 |     def read_qrels(self, *args, **kwargs) -> Dict[str, Dict[str, int]]:
272 |         qrels = defaultdict(dict)
273 |         for qrel in self.dataset.qrels_iter():
274 |             qrels[qrel.query_id][qrel.doc_id] = qrel.relevance
275 |         return qrels
276 | 
277 | 
278 | datasets_cls = {
279 |     "trec-covid": TrecCovid,
280 |     "trec-news": TrecNews,
281 |     "robust04": Robust04,
282 |     "touche": Touche,
283 | }
284 | 


--------------------------------------------------------------------------------
/inc_rel/download.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import shutil
  3 | import tarfile
  4 | from urllib.parse import urlparse
  5 | 
  6 | import ir_datasets
  7 | import requests
  8 | from requests.auth import HTTPBasicAuth
  9 | from settings import (
 10 |     Settings,
 11 |     TRECCovidDatasetSettings,
 12 |     TRECNewsDatasetSettings,
 13 |     TRECRobustDatasetSettings,
 14 |     WebisTouche2020DatasetSettings,
 15 | )
 16 | from tqdm.auto import tqdm
 17 | 
 18 | UNSET_USERNAME = "<NONE>"
 19 | 
 20 | 
 21 | def download_file(
 22 |     url: str,
 23 |     file_path: str,
 24 |     username: str = None,
 25 |     password: str = None,
 26 |     pbar_desc: str = "",
 27 |     response_attr: str = "raw",
 28 | ) -> None:
 29 |     """Download a file from a URL to a file path."""
 30 |     if os.path.isdir(file_path):
 31 |         file_path = os.path.join(file_path, os.path.basename(urlparse(url).path))
 32 | 
 33 |     if not os.path.exists(file_path):
 34 |         request_kwargs = {}
 35 |         if all([username, password]):
 36 |             request_kwargs["auth"] = HTTPBasicAuth(username, password)
 37 | 
 38 |         response = requests.get(url, stream=True, **request_kwargs)
 39 |         response.raise_for_status()
 40 | 
 41 |         if response_attr == "raw":
 42 |             total_length = int(response.headers.get("Content-Length", 0))
 43 |             with tqdm.wrapattr(
 44 |                 getattr(response, response_attr),
 45 |                 "read",
 46 |                 total=total_length,
 47 |                 ncols=100,
 48 |                 desc=pbar_desc,
 49 |             ) as pbar:
 50 |                 with open(file_path, "wb") as file:
 51 |                     shutil.copyfileobj(pbar, file)
 52 |         elif response_attr == "content":
 53 |             with open(file_path, "wb") as file:
 54 |                 file.write(response.content)
 55 |         else:
 56 |             raise ValueError(f"Invalid response_attr: {response_attr}")
 57 | 
 58 |     else:
 59 |         print(f"File already exists: {file_path}")
 60 | 
 61 |     return file_path
 62 | 
 63 | 
 64 | def untar_gzip(file_path: str, output_dir: str) -> None:
 65 |     """Untar a gzip file."""
 66 |     with tarfile.open(file_path, "r:gz") as tar:
 67 |         for name in tar.getnames():
 68 |             if not os.path.exists(os.path.join(output_dir, name)):
 69 |                 tar.extract(name, path=output_dir)
 70 | 
 71 | 
 72 | def main() -> None:
 73 |     # Setup data dir
 74 |     settings = Settings()
 75 |     os.makedirs(settings.raw_path, exist_ok=True)
 76 | 
 77 |     # Download TREC-COVID
 78 |     trec_covid_settings = TRECCovidDatasetSettings()
 79 |     os.makedirs(trec_covid_settings.raw_path, exist_ok=True)
 80 |     out_file = download_file(
 81 |         trec_covid_settings.corpus_url,
 82 |         trec_covid_settings.raw_path,
 83 |         pbar_desc="trec-covid/corpus",
 84 |     )
 85 |     untar_gzip(out_file, trec_covid_settings.raw_path)
 86 |     download_file(
 87 |         trec_covid_settings.topics_url,
 88 |         trec_covid_settings.raw_path,
 89 |         response_attr="content",
 90 |         pbar_desc="trec-covid/topcis",
 91 |     )
 92 |     download_file(
 93 |         trec_covid_settings.qrels_url,
 94 |         trec_covid_settings.raw_path,
 95 |         response_attr="content",
 96 |         pbar_desc="trec-covid/qrels",
 97 |     )
 98 | 
 99 |     # Download Webis Touche 2020
100 |     webis_touche_2020_settings = WebisTouche2020DatasetSettings()
101 |     ir_datasets.load(webis_touche_2020_settings.ir_datasets_name)
102 | 
103 |     # Download TREC-Robust
104 |     trec_robust_settings = TRECRobustDatasetSettings()
105 |     if trec_robust_settings.username == UNSET_USERNAME:
106 |         print(
107 |             "Skipping TREC-Robust download. Get credentials from TREC for this dataset."
108 |         )
109 |     else:
110 |         os.makedirs(trec_robust_settings.raw_path, exist_ok=True)
111 |         out_file = download_file(
112 |             trec_robust_settings.corpus_disk4_url,
113 |             trec_robust_settings.raw_path,
114 |             trec_robust_settings.username,
115 |             trec_robust_settings.password.get_secret_value(),
116 |             pbar_desc="trec-robust/corpus-disk4",
117 |         )
118 |         untar_gzip(out_file, trec_robust_settings.raw_path)
119 | 
120 |         out_file = download_file(
121 |             trec_robust_settings.corpus_disk5_url,
122 |             trec_robust_settings.raw_path,
123 |             trec_robust_settings.username,
124 |             trec_robust_settings.password.get_secret_value(),
125 |             pbar_desc="trec-robust/corpus-disk5",
126 |         )
127 |         untar_gzip(out_file, trec_robust_settings.raw_path)
128 | 
129 |         # create symlinks to disk4 and disk5 in ir_datasets
130 |         ir_datasets_robust_home_path = os.path.join(
131 |             settings.ir_datasets_home, "trec-robust04", "trec45"
132 |         )
133 |         os.makedirs(ir_datasets_robust_home_path, exist_ok=True)
134 |         for disk, robust_dataset in [
135 |             ("TREC-Disk-4", "FR94"),
136 |             ("TREC-Disk-4", "FT"),
137 |             ("TREC-Disk-5", "FBIS"),
138 |             ("TREC-Disk-5", "LATIMES"),
139 |         ]:
140 |             os.symlink(
141 |                 os.path.join(trec_robust_settings.raw_path, disk, robust_dataset),
142 |                 os.path.join(ir_datasets_robust_home_path, robust_dataset),
143 |             )
144 | 
145 |     # Download TREC-News
146 |     trec_news_settings = TRECNewsDatasetSettings()
147 |     if trec_news_settings.username == UNSET_USERNAME:
148 |         print(
149 |             "Skipping TREC-News download. Get credentials from TREC for this dataset."
150 |         )
151 |     else:
152 |         os.makedirs(trec_news_settings.raw_path, exist_ok=True)
153 |         out_file = download_file(
154 |             trec_news_settings.corpus_url,
155 |             trec_news_settings.raw_path,
156 |             trec_news_settings.username,
157 |             trec_news_settings.password.get_secret_value(),
158 |             pbar_desc="trec-news/corpus",
159 |         )
160 |         untar_gzip(out_file, trec_news_settings.raw_path)
161 |         download_file(
162 |             trec_news_settings.topics_url,
163 |             trec_news_settings.raw_path,
164 |             response_attr="content",
165 |             pbar_desc="trec-news/topics",
166 |         )
167 |         download_file(
168 |             trec_news_settings.qrels_url,
169 |             trec_news_settings.raw_path,
170 |             response_attr="content",
171 |             pbar_desc="trec-news/qrels",
172 |         )
173 | 
174 | 
175 | if __name__ == "__main__":
176 |     main()
177 | 


--------------------------------------------------------------------------------
/inc_rel/eval.py:
--------------------------------------------------------------------------------
 1 | from collections import defaultdict
 2 | from typing import Dict, List
 3 | 
 4 | import numpy as np
 5 | from pytrec_eval import RelevanceEvaluator
 6 | 
 7 | at_k = [1, 5, 10, 20, 50, 100, 1000]
 8 | map = "map_cut." + ",".join([str(k) for k in at_k])
 9 | ndcg = "ndcg_cut." + ",".join([str(k) for k in at_k])
10 | recall = "recall." + ",".join([str(k) for k in at_k])
11 | precision = "P." + ",".join([str(k) for k in at_k])
12 | rprecision = "Rprec." + ",".join([str(k) for k in at_k])
13 | measures = {map, ndcg, recall, precision, rprecision}
14 | 
15 | 
16 | def eval_bm25(
17 |     qrels: Dict[str, Dict[str, float]],
18 |     run: Dict[str, Dict[str, float]],
19 |     exclude_annotations: Dict[str, List[Dict]] = None,
20 | ) -> Dict[str, Dict[str, float]]:
21 | 
22 |     if exclude_annotations is not None:
23 |         filtered_qrels, filtered_run = {}, {}
24 |         for topic_id, annotation in exclude_annotations.items():
25 |             annotated_doc_ids = [a["doc_id"] for a in annotation]
26 | 
27 |             filtered_qrels[topic_id] = dict(
28 |                 filter(
29 |                     lambda doc_rel: doc_rel[0] not in annotated_doc_ids,
30 |                     qrels[topic_id].items(),
31 |                 )
32 |             )
33 |             filtered_run[topic_id] = dict(
34 |                 filter(
35 |                     lambda doc_rel: doc_rel[0] not in annotated_doc_ids,
36 |                     run[topic_id].items(),
37 |                 )
38 |             )
39 | 
40 |         run = filtered_run
41 |         qrels = filtered_qrels
42 | 
43 |     evaluator = RelevanceEvaluator(qrels, measures)
44 |     results = evaluator.evaluate(run)
45 | 
46 |     return results
47 | 
48 | 
49 | def accumulate_results(
50 |     bm25_eval: Dict[str, Dict[str, float]], topic_ids: List[str] = None
51 | ) -> Dict[str, float]:
52 |     acc_results = defaultdict(list)
53 |     for topic_id, measure2score in bm25_eval.items():
54 |         if topic_ids is not None and not topic_id in topic_ids:
55 |             continue
56 |         for measure, score in measure2score.items():
57 |             acc_results[measure].append(score)
58 |     mean = {k: np.mean(v) for k, v in acc_results.items()}
59 |     std = {k: np.std(v) for k, v in acc_results.items()}
60 | 
61 |     return {"mean": mean, "std": std}
62 | 


--------------------------------------------------------------------------------
/inc_rel/few_shot_trainer.py:
--------------------------------------------------------------------------------
  1 | import time
  2 | from typing import Any, Dict, List
  3 | 
  4 | import numpy as np
  5 | import torch
  6 | from reranking_evaluator import RerankingEvaluator
  7 | from sentence_transformers import (
  8 |     CrossEncoder,
  9 |     InputExample,
 10 |     SentenceTransformer,
 11 |     losses,
 12 | )
 13 | from torch.utils.data.dataloader import DataLoader
 14 | 
 15 | # from transformers import AdapterConfig
 16 | 
 17 | 
 18 | class FewShotTrainer:
 19 |     def __init__(
 20 |         self,
 21 |         model_name: str,
 22 |         ft_params: str,
 23 |         docs,
 24 |         initial_ranking,
 25 |         ranking_evaluator: RerankingEvaluator,
 26 |         pbar,
 27 |         model_path: str = None,
 28 |     ):
 29 | 
 30 |         self.model_name = model_name
 31 |         self.ft_params = ft_params
 32 |         self.docs = docs
 33 |         self.initial_ranking = initial_ranking
 34 |         self.ranking_evaluator = ranking_evaluator
 35 |         self.pbar = pbar
 36 |         self.model_path = model_path
 37 | 
 38 |     def init_model(self):
 39 |         if self.model_name.startswith("cross-encoder"):
 40 |             model = CrossEncoder(
 41 |                 self.model_name if self.model_path is None else self.model_path
 42 |             )
 43 |             architecture = "ce"
 44 |         elif self.model_name.startswith("sentence-transformer"):
 45 |             model = SentenceTransformer(
 46 |                 self.model_name if self.model_path is None else self.model_path
 47 |             )
 48 |             architecture = "bi"
 49 |         else:
 50 |             ValueError(self.model_name)
 51 | 
 52 |         return model, architecture
 53 | 
 54 |     def freeze_params(self):
 55 |         if self.architecture == "ce":
 56 |             np = self.model.model.named_parameters
 57 |         elif self.architecture == "bi":
 58 |             np = self.model.named_parameters
 59 |         else:
 60 |             raise ValueError(f"Unsupported Architecture: {self.architecture}")
 61 | 
 62 |         if self.ft_params == "head":
 63 |             if self.architecture == "bi":
 64 |                 raise RuntimeError("Head Tuning not compatiable with Bi-Encoder.")
 65 |             for name, param in np():
 66 |                 if "classifier" not in name:
 67 |                     param.requires_grad = False
 68 |         elif self.ft_params == "bias":
 69 |             for name, param in np():
 70 |                 if "bias" not in name:
 71 |                     param.requires_grad = False
 72 |         elif self.ft_params == "full":
 73 |             pass
 74 |         # elif self.ft_params == "adapter":
 75 |         #     if self.architecture == "bi":
 76 |         #         model_ref = self.model[0].auto_model
 77 |         #     elif self.architecture == "ce":
 78 |         #         model_ref = self.model.model
 79 |         #     adapter_config = AdapterConfig.load("pfeiffer", reduction_factor=16)
 80 |         #     adapter_name = "adapter"
 81 |         #     model_ref.add_adapter(adapter_name, config=adapter_config)
 82 | 
 83 |         #     model_ref.set_active_adapters("adapter")
 84 |         #     model_ref.train_adapter("adapter")
 85 |         else:
 86 |             raise ValueError(self.ft_params)
 87 | 
 88 |     def init_dataloader(self, annotations, batch_size):
 89 | 
 90 |         dataset = []
 91 |         for annotation in annotations:
 92 |             dataset.append(
 93 |                 InputExample(
 94 |                     texts=[annotation["query"], annotation["document"]],
 95 |                     label=float(annotation["label"] > 0),
 96 |                 )
 97 |             )
 98 |         g = torch.Generator()
 99 |         g.manual_seed(0)
100 |         return DataLoader(dataset, batch_size=batch_size, shuffle=True, generator=g)
101 | 
102 |     def train(
103 |         self,
104 |         topic2annotations,
105 |         epochs: int,
106 |         learning_rate: float,
107 |         update_pbar: bool = True,
108 |         iter_epochs: bool = True,
109 |         time_it: bool = False,
110 |     ) -> List[Dict[str, Any]]:
111 |         results = []
112 |         fit_times = []
113 |         eval_times = []
114 |         for topic_id, annotations in topic2annotations.items():
115 | 
116 |             # re-train the model for num_epochs each time from inital checkpoint
117 |             for num_epochs in range(1, epochs + 1):
118 |                 if not iter_epochs:
119 |                     if num_epochs != epochs:
120 |                         continue
121 | 
122 |                 torch.manual_seed(0)
123 |                 np.random.seed(0)
124 | 
125 |                 self.model, self.architecture = self.init_model()
126 |                 self.freeze_params()
127 | 
128 |                 t1 = time.time()
129 |                 self.fit(annotations, num_epochs, learning_rate)
130 |                 t2 = time.time()
131 |                 fit_times.append(t2 - t1)
132 | 
133 |                 t1 = time.time()
134 |                 metrics, run = self.eval(annotations)
135 |                 t2 = time.time()
136 |                 eval_times.append(t2 - t1)
137 | 
138 |                 results.append(
139 |                     {
140 |                         "topic_id": topic_id,
141 |                         "epoch": num_epochs,
142 |                         "learning_rate": learning_rate,
143 |                         "metrics": metrics,
144 |                         "run": run,
145 |                     }
146 |                 )
147 | 
148 |                 if update_pbar:
149 |                     self.pbar.update(1)
150 | 
151 |         if time_it:
152 |             return fit_times, eval_times
153 | 
154 |         return results
155 | 
156 |     def fit(self, annotations, num_epochs, learning_rate):
157 |         dataloader = self.init_dataloader(annotations, batch_size=1)
158 |         if self.architecture == "bi":
159 |             loss = losses.CosineSimilarityLoss(self.model)
160 |             train_args = {"train_objectives": [(dataloader, loss)]}
161 |             self.model.train()
162 |         elif self.architecture == "ce":
163 |             train_args = {"train_dataloader": dataloader}
164 |             self.model.model.train()
165 |         else:
166 |             raise ValueError(self.architecture)
167 | 
168 |         self.model.fit(
169 |             **train_args,
170 |             epochs=num_epochs,
171 |             optimizer_params={"lr": learning_rate},
172 |             save_best_model=False,
173 |             show_progress_bar=False,
174 |             warmup_steps=0,
175 |         )
176 | 
177 |     def eval(self, annotations):
178 | 
179 |         return self.ranking_evaluator(
180 |             model=self.model,
181 |             queries={annotations[0]["topic_id"]: annotations[0]["query"]},
182 |             docs=self.docs,
183 |             initial_ranking=self.initial_ranking,
184 |         )
185 | 


--------------------------------------------------------------------------------
/inc_rel/first_stage.py:
--------------------------------------------------------------------------------
 1 | import argparse
 2 | import json
 3 | import os
 4 | import time
 5 | 
 6 | from datasets import datasets_cls
 7 | from eval import accumulate_results, eval_bm25
 8 | from index import Index
 9 | from settings import dataset_settings_cls
10 | from tqdm.auto import tqdm
11 | 
12 | 
13 | def main(args):
14 |     index = Index()
15 | 
16 |     dataset = datasets_cls[args.name](
17 |         corpus_path=args.corpus_path,
18 |         topics_path=args.topics_path,
19 |         qrels_path=args.qrels_path,
20 |         load_corpus=not index.index_exists(index=args.name)
21 |         or args.name in ["trec-news"],
22 |         remove_topics_with_few_annotations=False,
23 |     )
24 | 
25 |     if not index.index_exists(index=args.name):
26 |         index.index_corpus(
27 |             index=args.name, corpus=dataset.corpus, corpus_size=dataset.corpus_size
28 |         )
29 | 
30 |     bm25_results, bm25_docs, _ = index.bm25_query(
31 |         index=args.name, topics=dataset.topics, size=args.bm25_size
32 |     )
33 |     bm25_eval = eval_bm25(dataset.qrels, bm25_results)
34 |     bm25_eval_acc = accumulate_results(bm25_eval)
35 | 
36 |     dataset_path = os.path.join(args.data_path, str(args.bm25_size))
37 |     os.makedirs(dataset_path, exist_ok=True)
38 |     for obj, name in zip(
39 |         [
40 |             bm25_results,
41 |             bm25_docs,
42 |             bm25_eval,
43 |             bm25_eval_acc,
44 |         ],
45 |         [
46 |             "full_bm25_results",
47 |             "full_bm25_docs",
48 |             "full_bm25_eval",
49 |             "full_bm25_eval_acc",
50 |         ],
51 |     ):
52 |         with open(os.path.join(dataset_path, f"{name}.json"), "w") as fh:
53 |             json.dump(obj, fh, indent=4)
54 | 
55 | 
56 | if __name__ == "__main__":
57 |     parser = argparse.ArgumentParser()
58 |     parser.add_argument("-d", "--dataset", type=str, required=True)
59 |     args = parser.parse_args()
60 |     print(args)
61 | 
62 |     dataset_settings = dataset_settings_cls[args.dataset]()
63 | 
64 |     main(dataset_settings)
65 | 


--------------------------------------------------------------------------------
/inc_rel/generate_few_shot.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import json
  3 | import os
  4 | import time
  5 | from collections import defaultdict
  6 | from pathlib import Path
  7 | 
  8 | import eval
  9 | import utils
 10 | from datasets import datasets_cls
 11 | from index import Index
 12 | from pydantic import SecretStr
 13 | from settings import dataset_settings_cls
 14 | from tqdm.auto import tqdm
 15 | 
 16 | 
 17 | def main(args):
 18 | 
 19 |     index = Index()
 20 | 
 21 |     dataset = datasets_cls[args.name](
 22 |         corpus_path=args.corpus_path,
 23 |         topics_path=args.topics_path,
 24 |         qrels_path=args.qrels_path,
 25 |         load_corpus=not index.index_exists(index=args.name)
 26 |         or args.name in ["trec-news"],
 27 |         remove_topics_with_few_negatives=args.enrich_bm25_path is None,
 28 |     )
 29 | 
 30 |     if not index.index_exists(index=args.name):
 31 |         index.index_corpus(
 32 |             index=args.name, corpus=dataset.corpus, corpus_size=dataset.corpus_size
 33 |         )
 34 | 
 35 |     # do first bm25 retrieval
 36 |     bm25_results, bm25_docs, _ = index.bm25_query(
 37 |         index=args.name, topics=dataset.topics, size=args.bm25_size
 38 |     )
 39 | 
 40 |     # we can use doc_id2text as corpus here because the bm25 docs are a superset of
 41 |     # the candidates
 42 |     max_k = max(args.num_samples)
 43 |     candidates = utils.sorted_annotated_retrieved_docs(
 44 |         corpus=bm25_docs,
 45 |         qrels=dataset.qrels,
 46 |         bm25_results=bm25_results,
 47 |         min_relevant_relevance=dataset.min_relevant_relevancy,
 48 |         remove_duplicates=args.remove_duplicates,
 49 |         enrich_negatives_from_bm25=args.enrich_bm25_path,
 50 |         num_non_relevant=max_k,
 51 |     )
 52 | 
 53 |     dataset_path = os.path.join(args.data_path, str(args.bm25_size))
 54 |     os.makedirs(dataset_path, exist_ok=True)
 55 | 
 56 |     # check if each topic has enough candidates in the bm25 results
 57 |     annotations = defaultdict(lambda: defaultdict(list))
 58 |     for topic_id, query in dataset.topics.copy().items():
 59 |         relevant_candidates = candidates[topic_id]["relevant"][:max_k]
 60 |         non_relevant_candidates = candidates[topic_id]["non_relevant"][:max_k]
 61 |         if not max_k == len(relevant_candidates) == len(non_relevant_candidates):
 62 | 
 63 |             print(
 64 |                 f"Could not find enough candidates for {topic_id=}. "
 65 |                 f"{len(relevant_candidates)=} {len(non_relevant_candidates)=} "
 66 |             )
 67 |             candidates.pop(topic_id)
 68 |             dataset.remove_topic(topic_id)
 69 | 
 70 |     # re-run bm25-query, since topics might have been removed
 71 |     bm25_results, bm25_docs, bm25_topics2time = index.bm25_query(
 72 |         index=args.name, topics=dataset.topics, size=args.bm25_size
 73 |     )
 74 |     bm25_eval = eval.eval_bm25(dataset.qrels, bm25_results)
 75 |     bm25_eval_acc = eval.accumulate_results(bm25_eval)
 76 | 
 77 |     # find annotations
 78 |     for topic_id, query in dataset.topics.items():
 79 |         relevant_candidates = candidates[topic_id]["relevant"][:max_k]
 80 |         non_relevant_candidates = candidates[topic_id]["non_relevant"][:max_k]
 81 |         for rel_doc_id, non_doc_id in zip(relevant_candidates, non_relevant_candidates):
 82 |             # if datsets has been enriched with bm25 negatives, qrels might not have
 83 |             # a doc, therefore we label it as 0
 84 |             rel_label = dataset.qrels[topic_id][rel_doc_id]
 85 |             rel_document = bm25_docs[rel_doc_id]
 86 | 
 87 |             non_label = dataset.qrels[topic_id].get(non_doc_id, 0)
 88 |             # incase of bm25 negatives, negatives might not be in bm_docs
 89 |             non_document = bm25_docs.get(
 90 |                 non_doc_id, index.get_doc_by_id(args.name, non_doc_id)
 91 |             )
 92 |             for k in args.num_samples:
 93 |                 # has to be 2k, since there are k relevant and k non-relevant
 94 |                 if len(annotations[k][topic_id]) < 2 * k:
 95 |                     for doc_id, label, document in zip(
 96 |                         [rel_doc_id, non_doc_id],
 97 |                         [rel_label, non_label],
 98 |                         [rel_document, non_document],
 99 |                     ):
100 |                         annotations[k][topic_id].append(
101 |                             {
102 |                                 "topic_id": topic_id,
103 |                                 "doc_id": doc_id,
104 |                                 "label": label,
105 |                                 "query": query,
106 |                                 "document": document,
107 |                             }
108 |                         )
109 |     for k, annotation in annotations.items():
110 |         with open(os.path.join(dataset_path, f"annotations_{k}.json"), "w") as fh:
111 |             json.dump(annotation, fh, indent=4)
112 | 
113 |     dataset.remove_annotations_from_qrels(annotations[max_k])
114 | 
115 |     # exclude annotations here, since bm25_results contain annotated docs
116 |     bm25_removed_eval = eval.eval_bm25(
117 |         dataset.qrels, bm25_results, exclude_annotations=annotations[max_k]
118 |     )
119 |     bm25_removed_eval_acc = eval.accumulate_results(bm25_removed_eval)
120 | 
121 |     for obj, name in zip(
122 |         [
123 |             dataset.topics,
124 |             dataset.qrels,
125 |             bm25_results,
126 |             bm25_docs,
127 |             bm25_eval,
128 |             bm25_eval_acc,
129 |             bm25_removed_eval,
130 |             bm25_removed_eval_acc,
131 |             bm25_topics2time,
132 |         ],
133 |         [
134 |             "topics",
135 |             "qrels",
136 |             "bm25_results",
137 |             "bm25_docs",
138 |             "bm25_eval",
139 |             "bm25_eval_accumulated",
140 |             "bm25_annot_removed_eval",
141 |             "bm25_annot_removed_eval_accumulated",
142 |             "bm25_topics2time",
143 |         ],
144 |     ):
145 |         with open(os.path.join(dataset_path, f"{name}.json"), "w") as fh:
146 |             json.dump(obj, fh, indent=4)
147 | 
148 |     for k in args.num_samples:
149 |         expan_results, expan_docs, expan_eval, expan_acc, expan_topic2time = (
150 |             {},
151 |             {},
152 |             {},
153 |             {},
154 |             {},
155 |         )
156 |         # expand bm25 query with annotated docs
157 |         bm25_key = "full"
158 |         (
159 |             expan_results[bm25_key],
160 |             expan_docs[bm25_key],
161 |             expan_topic2time[bm25_key],
162 |         ) = index.bm25_query_expansion(
163 |             index=args.name,
164 |             topics=dataset.topics,
165 |             annotations=annotations[k],
166 |             rm_annotations=annotations[max_k],
167 |             size=args.bm25_size,
168 |         )
169 |         expan_eval[bm25_key] = eval.eval_bm25(
170 |             qrels=dataset.qrels,
171 |             run=expan_results[bm25_key],
172 |         )
173 |         expan_acc[bm25_key] = eval.accumulate_results(expan_eval[bm25_key])
174 | 
175 |         for max_query_terms in [4, 8, 16, 32, 64]:
176 |             mlt_key = max_query_terms
177 |             (
178 |                 expan_results[mlt_key],
179 |                 expan_docs[mlt_key],
180 |                 expan_topic2time[mlt_key],
181 |             ) = index.more_like_this_bool_expansion(
182 |                 index=args.name,
183 |                 topics=dataset.topics,
184 |                 annotations=annotations[k],
185 |                 rm_annotations=annotations[max_k],
186 |                 size=args.bm25_size,
187 |                 max_query_terms=max_query_terms,
188 |             )
189 |             expan_eval[mlt_key] = eval.eval_bm25(
190 |                 qrels=dataset.qrels,
191 |                 run=expan_results[mlt_key],
192 |             )
193 |             expan_acc[mlt_key] = eval.accumulate_results(expan_eval[mlt_key])
194 | 
195 |         sample_path = os.path.join(args.data_path, str(args.bm25_size), f"k{k}")
196 |         os.makedirs(sample_path)
197 |         for key in expan_results.keys():
198 |             for name, obj in zip(
199 |                 ["results", "docs", "eval", "eval_acc", "times"],
200 |                 [expan_results, expan_docs, expan_eval, expan_acc, expan_topic2time],
201 |             ):
202 |                 with open(
203 |                     os.path.join(sample_path, f"expansion_{name}_{key}.json"), "w"
204 |                 ) as fh:
205 |                     json.dump(obj[key], fh, indent=4)
206 | 
207 |         for split_seed in args.split_seeds:
208 |             split_path = os.path.join(sample_path, f"s{str(split_seed)}")
209 |             os.makedirs(split_path)
210 | 
211 |             train_topics, valid_topics, test_topics = utils.create_split(
212 |                 topic_ids=list(annotations[k].keys()),
213 |                 seed=split_seed,
214 |                 split_sizes=args.split_sizes,
215 |             )
216 | 
217 |             splits = {}
218 |             for split_name, topic_ids in zip(
219 |                 ["train", "valid", "test"], [train_topics, valid_topics, test_topics]
220 |             ):
221 |                 splits[split_name] = {
222 |                     topic_id: v
223 |                     for topic_id, v in annotations[k].items()
224 |                     if topic_id in topic_ids
225 |                 }
226 |                 split_eval = {}
227 | 
228 |                 split_eval[("bm25", split_name)] = eval.accumulate_results(
229 |                     bm25_eval, list(splits[split_name].keys())
230 |                 )
231 |                 split_eval[("bm25-removed", split_name)] = eval.accumulate_results(
232 |                     bm25_removed_eval, list(splits[split_name].keys())
233 |                 )
234 | 
235 |                 for key in expan_results.keys():
236 |                     split_eval[(key, split_name)] = eval.accumulate_results(
237 |                         expan_eval[key], list(splits[split_name].keys())
238 |                     )
239 | 
240 |                 for split_name, split in splits.items():
241 |                     with open(
242 |                         os.path.join(split_path, f"{split_name}.json"), "w"
243 |                     ) as fh:
244 |                         json.dump(split, fh, indent=4)
245 | 
246 |                 for key in split_eval.keys():
247 |                     if isinstance(key[0], str):
248 |                         file_name = f"{split_name}_{key[0]}"
249 |                     else:
250 |                         file_name = f"{split_name}_expansion_{key[0]}"
251 |                     with open(
252 |                         os.path.join(split_path, f"{file_name}_eval_acc.json"), "w"
253 |                     ) as fh:
254 |                         json.dump(split_eval[key], fh, indent=4)
255 | 
256 |             config = vars(args)
257 |             config["k"] = k
258 |             config["seed"] = split_seed
259 |             config["min_annotations"] = dataset.min_annotations
260 |             config["min_relevant_relevancy"] = dataset.min_relevant_relevancy
261 |             config = {
262 |                 config_k: str(v) if isinstance(v, Path) else v
263 |                 for config_k, v in config.items()
264 |                 if not isinstance(v, SecretStr)
265 |             }
266 |             with open(os.path.join(split_path, "config.json"), "w") as fh:
267 |                 json.dump(config, fh, indent=4)
268 | 
269 | 
270 | if __name__ == "__main__":
271 |     parser = argparse.ArgumentParser()
272 |     parser.add_argument("-d", "--dataset", type=str, required=True)
273 |     args = parser.parse_args()
274 |     print(args)
275 | 
276 |     dataset_settings = dataset_settings_cls[args.dataset]()
277 | 
278 |     main(dataset_settings)
279 | 


--------------------------------------------------------------------------------
/inc_rel/index.py:
--------------------------------------------------------------------------------
  1 | import time
  2 | from typing import Callable, Dict, List, Union
  3 | 
  4 | import numpy as np
  5 | from elasticsearch import Elasticsearch
  6 | from elasticsearch import helpers as es_helpers
  7 | from tqdm.auto import tqdm
  8 | 
  9 | 
 10 | class Index:
 11 |     def __init__(self) -> None:
 12 | 
 13 |         self.es = Elasticsearch(["localhost"], timeout=3600, retry_after_timeout=True)
 14 | 
 15 |     def index_exists(self, index: str) -> bool:
 16 |         return self.es.indices.exists(index)
 17 | 
 18 |     def create_index(self, index: str):
 19 |         if self.index_exists(index):
 20 |             raise RuntimeError("Index already exists.")
 21 |         self.es.indices.create(
 22 |             index,
 23 |             body={
 24 |                 "mappings": {
 25 |                     "properties": {"text": {"type": "text", "term_vector": "yes"}}
 26 |                 }
 27 |             },
 28 |         )
 29 | 
 30 |     def index_corpus(
 31 |         self,
 32 |         index: str,
 33 |         corpus: Union[Callable, Dict[str, str]],
 34 |         corpus_size: Union[float, None] = None,
 35 |     ):
 36 |         self.create_index(index)
 37 |         chunk_size = 1024
 38 |         with tqdm(desc="indexing", ncols=100, total=corpus_size) as pbar:
 39 |             if isinstance(corpus, Callable):
 40 |                 chunk = []
 41 |                 for doc in corpus(verbose=False):
 42 |                     chunk.append(
 43 |                         {
 44 |                             "_index": index,
 45 |                             "_id": doc["id"],
 46 |                             "_source": {"text": doc["text"]},
 47 |                         }
 48 |                     )
 49 |                     if len(chunk) == chunk_size:
 50 |                         es_helpers.bulk(self.es, chunk)
 51 |                         pbar.update(len(chunk))
 52 |                         chunk = []
 53 |                 es_helpers.bulk(self.es, chunk)
 54 |                 pbar.update(len(chunk))
 55 |                 time.sleep(0.1)
 56 |             else:
 57 |                 for start_idx in range(0, len(corpus), chunk_size):
 58 |                     end_idx = start_idx + chunk_size
 59 |                     chunk = [
 60 |                         {
 61 |                             "_index": index,
 62 |                             "_id": doc_id,
 63 |                             "_source": {"text": corpus[doc_id]},
 64 |                         }
 65 |                         for doc_id in list(corpus.keys())[start_idx:end_idx]
 66 |                     ]
 67 |                     es_helpers.bulk(self.es, chunk)
 68 |                     pbar.update(chunk_size)
 69 |                     time.sleep(0.1)
 70 | 
 71 |         for _ in tqdm(range(120), ncols=100, desc="Granting ES some beauty sleep."):
 72 |             time.sleep(1)
 73 | 
 74 |     def get_doc_by_id(self, index: str, doc_id: str):
 75 |         return self.es.get(index=index, id=doc_id)["_source"]["text"]
 76 | 
 77 |     def query(self, index, body, time_it: bool = False):
 78 |         if time_it:
 79 |             self.es.indices.clear_cache(
 80 |                 index=index,
 81 |             )
 82 |             time.sleep(1)
 83 |         result = self.es.search(index=index, body=body)
 84 |         return result
 85 | 
 86 |     def filter_labeled_docs(
 87 |         self, remove_topics_from_annotatoins: List[Dict], result: List, size: int
 88 |     ):
 89 |         labeled_doc_ids = list(
 90 |             map(
 91 |                 lambda a: a["doc_id"],
 92 |                 remove_topics_from_annotatoins,
 93 |             )
 94 |         )
 95 |         result = [r for r in result if r["_id"] not in labeled_doc_ids][:size]
 96 |         return result
 97 | 
 98 |     def bm25_query(
 99 |         self, index: str, topics: Dict[str, str], size: int, time_it: bool = False
100 |     ):
101 |         bm25_results = {}
102 |         es_query = {
103 |             "query": {
104 |                 "multi_match": {
105 |                     "query": None,
106 |                     "fields": ["text"],
107 |                 }
108 |             },
109 |             "size": size,
110 |         }
111 | 
112 |         bm25_results, doc_id2text, topic_id2took = {}, {}, {}
113 |         for topic_id, query in tqdm(
114 |             topics.items(), total=len(topics), desc="caching bm25", ncols=100
115 |         ):
116 |             es_query["query"]["multi_match"]["query"] = query
117 |             result = self.query(index, body=es_query, time_it=time_it)
118 |             hits = result["hits"]["hits"]
119 |             topic_id2took[topic_id] = result["took"]
120 | 
121 |             doc_id2score = {}
122 |             for r in hits:
123 |                 doc_id2score[r["_id"]] = r["_score"]
124 |                 doc_id2text[r["_id"]] = r["_source"]["text"]
125 |             bm25_results[topic_id] = doc_id2score
126 | 
127 |         times = list(topic_id2took.values())
128 |         topic_id2took["avg"] = np.mean(times)
129 |         topic_id2took["std"] = np.std(times)
130 | 
131 |         return bm25_results, doc_id2text, topic_id2took
132 | 
133 |     def more_like_this_expansion(
134 |         self,
135 |         index: str,
136 |         topics: Dict[str, str],
137 |         annotations: Dict[str, List[Dict]],
138 |         rm_annotations: Dict[str, List[Dict]],
139 |         size: int,
140 |         unlike: bool = False,
141 |         max_query_terms: int = 25,
142 |         time_it: bool = False,
143 |     ):
144 |         # https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-mlt-query.html#query-dsl-mlt-query
145 |         es_query = {
146 |             "query": {
147 |                 "more_like_this": {
148 |                     "fields": ["text"],
149 |                     "max_query_terms": max_query_terms,
150 |                 }
151 |             },
152 |             "size": size + len(rm_annotations[list(rm_annotations.keys())[0]]),
153 |         }
154 |         mlt_results, doc_id2text, topic_id2took = {}, {}, {}
155 |         for topic_id, annotation in tqdm(
156 |             annotations.items(),
157 |             total=len(annotations),
158 |             desc=f"caching MLT expansion (terms={max_query_terms:04d}, negatives={unlike})",
159 |             ncols=100,
160 |         ):
161 | 
162 |             query = topics[topic_id]
163 |             like_query, unlike_query = (
164 |                 [],
165 |                 [],
166 |             )
167 |             for a in annotation:
168 |                 if a["label"] >= 1:
169 |                     like_query.append({"_index": index, "_id": a["doc_id"]})
170 |                 else:
171 |                     unlike_query.append({"_index": index, "_id": a["doc_id"]})
172 | 
173 |             es_query["query"]["more_like_this"]["like"] = like_query + [query]
174 |             if unlike:
175 |                 es_query["query"]["more_like_this"]["unlike"] = unlike_query
176 | 
177 |             result = self.query(index, body=es_query, time_it=time_it)
178 |             hits = result["hits"]["hits"]
179 |             topic_id2took[topic_id] = result["took"]
180 |             hits = self.filter_labeled_docs(rm_annotations[topic_id], hits, size)
181 |             doc_id2score = {}
182 |             for r in hits:
183 |                 doc_id2score[r["_id"]] = r["_score"]
184 |                 doc_id2text[r["_id"]] = r["_source"]["text"]
185 |             mlt_results[topic_id] = doc_id2score
186 | 
187 |         times = list(topic_id2took.values())
188 |         topic_id2took["avg"] = np.mean(times)
189 |         topic_id2took["std"] = np.std(times)
190 | 
191 |         return mlt_results, doc_id2text, topic_id2took
192 | 
193 |     def bm25_query_expansion(
194 |         self,
195 |         index: str,
196 |         topics: Dict[str, str],
197 |         annotations: Dict[str, List[Dict]],
198 |         rm_annotations: Dict[str, List[Dict]],
199 |         size: int,
200 |         time_it: bool = False,
201 |     ):
202 |         es_query = {
203 |             "query": {
204 |                 "bool": {
205 |                     "must": None,
206 |                     "should": None,
207 |                 }
208 |             },
209 |             "size": size + len(rm_annotations[list(rm_annotations.keys())[0]]),
210 |         }
211 | 
212 |         bm25_results, doc_id2text, topic_id2took = {}, {}, {}
213 |         for topic_id, annotation in tqdm(
214 |             annotations.items(),
215 |             total=len(annotations),
216 |             desc="caching bm25 expansion",
217 |             ncols=100,
218 |         ):
219 | 
220 |             query = topics[topic_id]
221 | 
222 |             relevant_docs = []
223 |             for a in annotation:
224 |                 if a["label"] >= 1:
225 |                     relevant_docs.append(a["document"])
226 |             # add query
227 |             es_query["query"]["bool"]["must"] = [
228 |                 {
229 |                     "match": {
230 |                         "text": {
231 |                             "query": query,
232 |                         },
233 |                     }
234 |                 },
235 |             ]
236 |             # add relevant docs
237 |             es_query["query"]["bool"]["should"] = [
238 |                 {
239 |                     "match": {
240 |                         "text": {
241 |                             "query": doc,
242 |                             "boost": 1,
243 |                         },
244 |                     }
245 |                 }
246 |                 for doc in relevant_docs
247 |             ]
248 |             result = self.query(index, body=es_query, time_it=time_it)
249 |             hits = result["hits"]["hits"]
250 |             topic_id2took[topic_id] = result["took"]
251 |             hits = self.filter_labeled_docs(rm_annotations[topic_id], hits, size)
252 |             doc_id2score = {}
253 |             for r in hits:
254 |                 doc_id2score[r["_id"]] = r["_score"]
255 |                 doc_id2text[r["_id"]] = r["_source"]["text"]
256 |             bm25_results[topic_id] = doc_id2score
257 | 
258 |         times = list(topic_id2took.values())
259 |         topic_id2took["avg"] = np.mean(times)
260 |         topic_id2took["std"] = np.std(times)
261 | 
262 |         return bm25_results, doc_id2text, topic_id2took
263 | 
264 |     def more_like_this_bool_expansion(
265 |         self,
266 |         index: str,
267 |         topics: Dict[str, str],
268 |         annotations: Dict[str, List[Dict]],
269 |         rm_annotations: Dict[str, List[Dict]],
270 |         size: int,
271 |         max_query_terms: int = 25,
272 |         time_it: bool = False,
273 |     ):
274 |         # https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-mlt-query.html#query-dsl-mlt-query
275 |         def get_mlt_query(max_query_terms=25):
276 |             return {
277 |                 "more_like_this": {
278 |                     "fields": ["text"],
279 |                     "max_query_terms": max_query_terms,
280 |                 }
281 |             }
282 | 
283 |         def get_query(size):
284 |             return {
285 |                 "query": {
286 |                     "bool": {
287 |                         "must": None,
288 |                         "should": [],
289 |                     }
290 |                 },
291 |                 "size": size + len(rm_annotations[list(rm_annotations.keys())[0]]),
292 |             }
293 | 
294 |         mlt_results, doc_id2text, topic_id2took = {}, {}, {}
295 |         for topic_id, annotation in tqdm(
296 |             annotations.items(),
297 |             total=len(annotations),
298 |             desc=f"caching MLT expansion (terms={max_query_terms:04d})",
299 |             ncols=100,
300 |         ):
301 | 
302 |             query = topics[topic_id]
303 |             es_query = get_query(size)
304 |             es_query["query"]["bool"]["must"] = {
305 |                 "match": {
306 |                     "text": {
307 |                         "query": query,
308 |                     },
309 |                 }
310 |             }
311 |             for a in annotation:
312 |                 if a["label"] >= 1:
313 |                     es_query["query"]["bool"]["should"].append(
314 |                         get_mlt_query(max_query_terms)
315 |                     )
316 |                     es_query["query"]["bool"]["should"][-1]["more_like_this"][
317 |                         "like"
318 |                     ] = {"_index": index, "_id": a["doc_id"]}
319 | 
320 |             result = self.query(index, body=es_query, time_it=time_it)
321 |             hits = result["hits"]["hits"]
322 |             topic_id2took[topic_id] = result["took"]
323 |             hits = self.filter_labeled_docs(rm_annotations[topic_id], hits, size)
324 |             doc_id2score = {}
325 |             for r in hits:
326 |                 doc_id2score[r["_id"]] = r["_score"]
327 |                 doc_id2text[r["_id"]] = r["_source"]["text"]
328 |             mlt_results[topic_id] = doc_id2score
329 | 
330 |         times = list(topic_id2took.values())
331 |         topic_id2took["avg"] = np.mean(times)
332 |         topic_id2took["std"] = np.std(times)
333 | 
334 |         return mlt_results, doc_id2text, topic_id2took
335 | 


--------------------------------------------------------------------------------
/inc_rel/knn_eval.py:
--------------------------------------------------------------------------------
  1 | import ast
  2 | import json
  3 | import operator
  4 | import os
  5 | from collections import defaultdict
  6 | from typing import Dict, List
  7 | 
  8 | import eval
  9 | import simple_parsing
 10 | from args import KNN
 11 | from reranking_evaluator import RerankingEvaluator
 12 | 
 13 | 
 14 | def collect_similarities(
 15 |     qrels,
 16 |     annotations,
 17 |     similarities,
 18 |     expansion_results,
 19 |     query_sim: bool,
 20 |     annot_sim: bool,
 21 |     relevant_or_non_relevant: str,
 22 | ):
 23 |     sims = defaultdict(lambda: defaultdict(list))
 24 |     if relevant_or_non_relevant == "relevant":
 25 |         op = operator.gt
 26 |     elif relevant_or_non_relevant == "non-relevant":
 27 |         op = operator.le
 28 |     else:
 29 |         raise ValueError
 30 |     for topic_id in qrels.keys():
 31 |         if query_sim:
 32 |             # add query:doc similarity score
 33 |             for doc_id in expansion_results[topic_id].keys():
 34 |                 sims[topic_id][doc_id].append(similarities[(topic_id, doc_id)])
 35 |         if annot_sim:
 36 |             # add relevant-annotation:doc similarity score
 37 |             for annotation in annotations[topic_id]:
 38 |                 if op(annotation["label"], 0):
 39 |                     for doc_id in expansion_results[topic_id].keys():
 40 |                         sims[topic_id][doc_id].append(
 41 |                             similarities[(annotation["doc_id"], doc_id)]
 42 |                         )
 43 |     return sims
 44 | 
 45 | 
 46 | def aggregate_sim_collection(sims, fn):
 47 |     run = defaultdict(dict)
 48 |     for topic_id, docids2scores in sims.items():
 49 |         for doc_id, scores in docids2scores.items():
 50 |             run[topic_id][doc_id] = fn(scores)
 51 |     return run
 52 | 
 53 | 
 54 | def main(args):
 55 |     similarities_path = os.path.join(args.exp_path, f"similarities.json")
 56 |     with open(similarities_path) as fh:
 57 |         similarities = json.load(fh)
 58 |     similarities = {ast.literal_eval(k): v for k, v in similarities.items()}
 59 | 
 60 |     for k in args.num_samples:
 61 |         with open(
 62 |             os.path.join(args.data_path, f"k{k}", f"expansion_results_16.json")
 63 |         ) as fh:
 64 |             expansion_results: Dict[str, Dict] = json.load(fh)
 65 | 
 66 |         with open(os.path.join(args.data_path, f"annotations_{k}.json")) as fh:
 67 |             annotations: Dict[str, List[Dict]] = json.load(fh)
 68 | 
 69 |         query_only_sims = collect_similarities(
 70 |             args.qrels,
 71 |             annotations,
 72 |             similarities,
 73 |             expansion_results,
 74 |             query_sim=True,
 75 |             annot_sim=False,
 76 |             relevant_or_non_relevant="relevant",
 77 |         )
 78 |         annot_only_sims = collect_similarities(
 79 |             args.qrels,
 80 |             annotations,
 81 |             similarities,
 82 |             expansion_results,
 83 |             query_sim=False,
 84 |             annot_sim=True,
 85 |             relevant_or_non_relevant="relevant",
 86 |         )
 87 |         query_annot_sims = collect_similarities(
 88 |             args.qrels,
 89 |             annotations,
 90 |             similarities,
 91 |             expansion_results,
 92 |             query_sim=True,
 93 |             annot_sim=True,
 94 |             relevant_or_non_relevant="relevant",
 95 |         )
 96 | 
 97 |         results = {}
 98 |         full_result = defaultdict(dict)
 99 |         for sim, sim_name in zip(
100 |             [query_only_sims, annot_only_sims, query_annot_sims],
101 |             ["query", "annot", "query_annot"],
102 |         ):
103 |             for fn, fn_name in zip([sum, max], ["sum", "max"]):
104 |                 sims_agg = aggregate_sim_collection(sim, fn=fn)
105 |                 results[(k, sim_name, fn_name)] = RerankingEvaluator(args.qrels).eval(
106 |                     sims_agg
107 |                 )
108 | 
109 |                 full_result[(k, sim_name, fn_name)] = [
110 |                     {
111 |                         "topic_id": topic_id,
112 |                         "metrics": results[(k, sim_name, fn_name)][topic_id],
113 |                         "run": {
114 |                             topic_id: dict(
115 |                                 sorted(
116 |                                     sims_agg[topic_id].items(),
117 |                                     key=lambda item: item[1],
118 |                                     reverse=True,
119 |                                 )
120 |                             )
121 |                         },
122 |                     }
123 |                     for topic_id in sim.keys()
124 |                 ]
125 | 
126 |         annot_non_rel_sims = collect_similarities(
127 |             args.qrels,
128 |             annotations,
129 |             similarities,
130 |             expansion_results,
131 |             query_sim=False,
132 |             annot_sim=True,
133 |             relevant_or_non_relevant="non-relevant",
134 |         )
135 |         annot_non_rel_agg = aggregate_sim_collection(sim, fn=sum)
136 |         sim_name = "annot_non_rel"
137 |         for fn, fn_name in zip([sum, max], ["sum", "max"]):
138 |             sims_agg = aggregate_sim_collection(sim, fn=fn)
139 |             for topic_id in annot_non_rel_agg.keys():
140 |                 for doc_id in annot_non_rel_agg[topic_id].keys():
141 |                     sims_agg[topic_id][doc_id] -= annot_non_rel_agg[topic_id][doc_id]
142 |             results[(k, sim_name, fn_name)] = RerankingEvaluator(args.qrels).eval(
143 |                 sims_agg
144 |             )
145 |             full_result[(k, sim_name, fn_name)] = [
146 |                 {
147 |                     "topic_id": topic_id,
148 |                     "metrics": results[(k, sim_name, fn_name)][topic_id],
149 |                     "run": {
150 |                         topic_id: dict(
151 |                             sorted(
152 |                                 sims_agg[topic_id].items(),
153 |                                 key=lambda item: item[1],
154 |                                 reverse=True,
155 |                             )
156 |                         )
157 |                     },
158 |                 }
159 |                 for topic_id in sim.keys()
160 |             ]
161 | 
162 |         for key in full_result.keys():
163 |             _k, sim_name, fn_name = key
164 |             eval_file = os.path.join(
165 |                 args.exp_path, f"k{_k}_{sim_name}-{fn_name}_eval.json"
166 |             )
167 |             with open(eval_file, "w") as fh:
168 |                 json.dump(full_result[key], fh, indent=4)
169 | 
170 |         split2metric = defaultdict(list)
171 |         for seed in args.seeds:
172 |             for split in args.splits:
173 |                 with open(
174 |                     os.path.join(args.data_path, f"k{k}", f"s{seed}", f"{split}.json")
175 |                 ) as fh:
176 |                     split_seed = json.load(fh)
177 | 
178 |                 for exp_name, result in results.items():
179 |                     split_seed_eval_acc = eval.accumulate_results(
180 |                         result, topic_ids=list(split_seed.keys())
181 |                     )
182 | 
183 |                     _, sim_name, fn_name = exp_name
184 | 
185 |                     eval_acc_file = os.path.join(
186 |                         args.exp_path,
187 |                         f"k{k}_s{seed}_{split}_{sim_name}-{fn_name}_eval_acc.json",
188 |                     )
189 |                     with open(eval_acc_file, "w") as fh:
190 |                         json.dump(split_seed_eval_acc, fh, indent=4)
191 |                     if sim_name == "query_annot" and fn_name == "sum":
192 |                         # works best
193 |                         split2metric[split].append(
194 |                             split_seed_eval_acc["mean"][args.metric]
195 |                         )
196 |                         print(
197 |                             f"k={k:02d} split={split:5s} seed={seed:02d} "
198 |                             f"{args.metric}={split_seed_eval_acc['mean'][args.metric]:.4f}"
199 |                         )
200 | 
201 |         print("---MEAN---")
202 |         for split in args.splits:
203 |             print(
204 |                 f"k={k:02d} split={split:5s} "
205 |                 f"{args.metric}={sum(split2metric[split]) / len(split2metric[split]):.4f}"
206 |             )
207 | 
208 | 
209 | if __name__ == "__main__":
210 |     args = simple_parsing.parse(KNN)
211 |     print(args)
212 |     main(args)
213 | 


--------------------------------------------------------------------------------
/inc_rel/knn_index.py:
--------------------------------------------------------------------------------
 1 | import json
 2 | import os
 3 | from typing import Dict, List
 4 | 
 5 | import simple_parsing
 6 | import torch
 7 | from args import KNN
 8 | from sentence_transformers import SentenceTransformer
 9 | 
10 | 
11 | def main(args):
12 | 
13 |     docs = {}
14 |     for k in args.num_samples:
15 |         base_path = os.path.join(args.data_path, f"k{k}")
16 |         docs_file = os.path.join(base_path, f"expansion_docs_16.json")
17 |         with open(docs_file) as fh:
18 |             docs.update(json.load(fh))
19 | 
20 |     with open(os.path.join(args.data_path, "topics.json")) as fh:
21 |         queries = json.load(fh)
22 | 
23 |     with open(os.path.join(args.data_path, "annotations_8.json")) as fh:
24 |         annotations: Dict[List[Dict]] = json.load(fh)
25 |     annotations = {
26 |         a["doc_id"]: a["document"]
27 |         for annotation in annotations.values()
28 |         for a in annotation
29 |     }
30 | 
31 |     print(
32 |         f"Loaded {len(queries)} queries "
33 |         f"{len(annotations)} annotation docs "
34 |         f"and {len(docs)} retrieval docs."
35 |     )
36 | 
37 |     model = SentenceTransformer(args.model)
38 |     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
39 |     model.to(device)
40 |     model.eval()
41 | 
42 |     with torch.no_grad():
43 |         for name, id2doc in zip(
44 |             ["queries", "annotations", "docs"], [queries, annotations, docs]
45 |         ):
46 |             embeddings_path = os.path.join(args.exp_path, f"{name}_embeddings.json")
47 |             if os.path.exists(embeddings_path):
48 |                 print(f"{name} embeddings already exists.")
49 |             else:
50 |                 texts = list(id2doc.values())
51 |                 texts_embeddings = model.encode(texts, show_progress_bar=True)
52 |                 texts_embeddings_mapped = {}
53 |                 for i, _id in enumerate(id2doc.keys()):
54 |                     texts_embeddings_mapped[_id] = texts_embeddings[i].tolist()
55 |                 with open(embeddings_path, "w") as fh:
56 |                     json.dump(texts_embeddings_mapped, fh, indent=4)
57 | 
58 | 
59 | if __name__ == "__main__":
60 |     args = simple_parsing.parse(KNN)
61 |     print(args)
62 |     main(args)
63 | 


--------------------------------------------------------------------------------
/inc_rel/knn_similarities.py:
--------------------------------------------------------------------------------
 1 | import json
 2 | import os
 3 | from typing import Dict, List
 4 | 
 5 | import numpy as np
 6 | import simple_parsing
 7 | from args import KNN
 8 | from sentence_transformers import util
 9 | from tqdm.auto import tqdm
10 | 
11 | 
12 | def main(args):
13 | 
14 |     if args.scoring_fn == "cos":
15 |         scoring_fn = util.cos_sim
16 |     elif args.scoring_fn == "dot":
17 |         scoring_fn = util.dot_score
18 | 
19 |     embeddings = {}
20 |     for name in ["queries", "annotations", "docs"]:
21 |         embeddings_path = os.path.join(args.exp_path, f"{name}_embeddings.json")
22 |         with open(embeddings_path) as fh:
23 |             embeddings[name] = json.load(fh)
24 | 
25 |     with open(os.path.join(args.data_path, "annotations_8.json")) as fh:
26 |         annotations: Dict[Dict[List]] = json.load(fh)
27 | 
28 |     similarites = {}
29 |     for k in [2, 4, 8]:
30 |         expansion_results_path = os.path.join(
31 |             args.data_path, f"k{k}", f"expansion_results_16.json"
32 |         )
33 |         with open(expansion_results_path) as fh:
34 |             expansion_results: Dict[Dict] = json.load(fh)
35 |         for topic_id, doc_id2score in tqdm(
36 |             expansion_results.items(),
37 |             total=len(expansion_results),
38 |             ncols=100,
39 |             desc="expansion_results",
40 |         ):
41 |             query_embedding = np.asarray(embeddings["queries"][topic_id]).reshape(1, -1)
42 |             doc_ids = doc_id2score.keys()
43 |             doc_embeddings = []
44 |             new_keys = []
45 |             for doc_id in doc_ids:
46 |                 key = (topic_id, doc_id)
47 |                 if key not in similarites:
48 |                     new_keys.append(key)
49 |                     doc_embeddings.append(embeddings["docs"][doc_id])
50 | 
51 |             if doc_embeddings:
52 |                 doc_embeddings = np.asarray(doc_embeddings)
53 |                 cos_sim = scoring_fn(query_embedding, doc_embeddings)
54 |                 for j, key in enumerate(new_keys):
55 |                     similarites[key] = cos_sim[0, j].item()
56 |             else:
57 |                 print(f"No new query-doc pairs found for topic={topic_id} at k={k}")
58 | 
59 |         for topic_id, annotation in tqdm(
60 |             annotations.items(), total=len(annotations), ncols=100, desc="Annotations"
61 |         ):
62 |             annotation_doc_ids = [a["doc_id"] for a in annotation]
63 |             annotation_embeddings = [
64 |                 embeddings["annotations"][doc_id] for doc_id in annotation_doc_ids
65 |             ]
66 | 
67 |             for i, annotation_doc_id in enumerate(annotation_doc_ids):
68 |                 doc_embeddings = []
69 |                 new_keys = []
70 |                 for doc_id in expansion_results[topic_id].keys():
71 |                     key = (annotation_doc_id, doc_id)
72 |                     if key not in similarites:
73 |                         new_keys.append(key)
74 |                         doc_embeddings.append(embeddings["docs"][doc_id])
75 |                 if doc_embeddings:
76 |                     doc_embeddings = np.asarray(doc_embeddings)
77 |                     annotation_embedding = np.asarray(annotation_embeddings[i]).reshape(
78 |                         1, -1
79 |                     )
80 |                     cos_sim = util.cos_sim(annotation_embedding, doc_embeddings)
81 |                     for j, key in enumerate(new_keys):
82 |                         similarites[key] = cos_sim[0, j].item()
83 |                 else:
84 |                     print(
85 |                         "No new annot-doc pairs found for "
86 |                         f"topic={topic_id} at k={k} for annot={i, annotation_doc_id}"
87 |                     )
88 | 
89 |     similarites_path = os.path.join(args.exp_path, f"similarities.json")
90 |     with open(similarites_path, "w") as fh:
91 |         json.dump({str(k): v for k, v in similarites.items()}, fh, indent=4)
92 | 
93 | 
94 | if __name__ == "__main__":
95 |     args = simple_parsing.parse(KNN)
96 |     print(args)
97 |     main(args)
98 | 


--------------------------------------------------------------------------------
/inc_rel/meta/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/UKPLab/incorporating-relevance/c7419c136f0b07bfe105095ebb49d2bb088c00b6/inc_rel/meta/__init__.py


--------------------------------------------------------------------------------
/inc_rel/meta/adapter_meta.py:
--------------------------------------------------------------------------------
 1 | import logging
 2 | from typing import Dict
 3 | 
 4 | import torch.nn as nn
 5 | import torch.nn.functional as F
 6 | from meta.module_meta import ModuleMetaMixin
 7 | 
 8 | logger = logging.getLogger(__name__)
 9 | 
10 | 
11 | class AdapterMetaMixin(ModuleMetaMixin):
12 |     @property
13 |     def layer(self):
14 |         return self._layer
15 | 
16 |     @layer.setter
17 |     def layer(self, layer):
18 |         self._layer = layer
19 | 
20 |     def get_params(self):
21 |         return {
22 |             f"layer.{self.layer}.adapter_down.weight": self.adapter_down[0].weight,
23 |             f"layer.{self.layer}.adapter_down.bias": self.adapter_down[0].bias,
24 |             f"layer.{self.layer}.adapter_up.weight": self.adapter_up.weight,
25 |             f"layer.{self.layer}.adapter_up.bias": self.adapter_up.bias,
26 |         }
27 | 
28 |     def set_params(self, params):
29 |         params = dict(filter(lambda kv: f"layer.{self.layer}" in kv[0], params.items()))
30 |         self.params = params
31 | 
32 |     def forward(self, x, residual_input):
33 |         """forward pass with functional functions for meta learning"""
34 | 
35 |         down = F.linear(
36 |             input=x,
37 |             weight=self.params[f"layer.{self.layer}.adapter_down.weight"],
38 |             bias=self.params[f"layer.{self.layer}.adapter_down.bias"],
39 |         )
40 |         down = self.adapter_down[1](down)  # activation function
41 | 
42 |         up = F.linear(
43 |             input=down,
44 |             weight=self.params[f"layer.{self.layer}.adapter_up.weight"],
45 |             bias=self.params[f"layer.{self.layer}.adapter_up.bias"],
46 |         )
47 | 
48 |         output = up
49 | 
50 |         # apply residual connection before layer norm if configured in this way
51 |         if self.residual_before_ln:
52 |             output = output + residual_input
53 | 
54 |         # apply layer norm if available
55 |         if self.add_layer_norm_after:
56 |             output = self.adapter_norm_after(
57 |                 output
58 |             )  # this can stay since its not affected by meta learning
59 | 
60 |         # if residual should be applied after layer norm, apply it here
61 |         if not self.residual_before_ln:
62 |             output = output + residual_input
63 | 
64 |         return output, down, up
65 | 
66 | 
67 | def extend_text_adapter(model):
68 |     for layer_i, layer in enumerate(model.encoder.layer):
69 |         for adapter_obj in layer.output.adapters.values():
70 |             # https://stackoverflow.com/a/31075641
71 |             base_cls = adapter_obj.__class__
72 |             base_cls_name = adapter_obj.__class__.__name__
73 |             adapter_obj.__class__ = type(
74 |                 base_cls_name, (AdapterMetaMixin, base_cls), {}
75 |             )
76 |             adapter_obj.layer = layer_i
77 | 


--------------------------------------------------------------------------------
/inc_rel/meta/bias_meta.py:
--------------------------------------------------------------------------------
 1 | import torch.nn as nn
 2 | import torch.nn.functional as F
 3 | from meta.module_meta import ModuleMetaMixin
 4 | 
 5 | 
 6 | class LinearBiasMetaMixin(ModuleMetaMixin):
 7 |     @property
 8 |     def module_name(self):
 9 |         return self._module_name
10 | 
11 |     @module_name.setter
12 |     def module_name(self, module_name):
13 |         self._module_name = module_name
14 | 
15 |     def get_params(self):
16 |         return {self.module_name: self.bias}
17 | 
18 |     def set_params(self, params):
19 |         self.params = {self.module_name: params[self.module_name]}
20 | 
21 |     def forward(self, x):
22 |         """forward pass with functional functions for meta learning"""
23 | 
24 |         return F.linear(
25 |             input=x,
26 |             weight=self.weight,
27 |             bias=self.params[self.module_name],
28 |         )
29 | 
30 | 
31 | def extend_linear(model, prefix=""):
32 | 
33 |     for n, m in model.encoder.named_modules(prefix=prefix):
34 |         if isinstance(m, nn.Linear):
35 |             base_cls = m.__class__
36 |             m.__class__ = type("MetaLinear", (LinearBiasMetaMixin, base_cls), {})
37 |             m.module_name = n + ".bias"
38 | 


--------------------------------------------------------------------------------
/inc_rel/meta/classification_head_meta.py:
--------------------------------------------------------------------------------
 1 | import logging
 2 | from typing import Dict
 3 | 
 4 | import torch.nn as nn
 5 | import torch.nn.functional as F
 6 | from meta.module_meta import ModuleMetaMixin
 7 | 
 8 | logger = logging.getLogger(__name__)
 9 | 
10 | 
11 | class ClassificationHeadMixin(ModuleMetaMixin):
12 |     def get_params(self):
13 |         params = {
14 |             "clf.linear.weight": self.weight,
15 |             "clf.linear.bias": self.bias,
16 |         }
17 |         return params
18 | 
19 |     def set_params(self, params):
20 |         params = dict(filter(lambda kv: f"clf" in kv[0], params.items()))
21 |         self.params = params
22 | 
23 |     def forward(self, input):
24 | 
25 |         return F.linear(
26 |             input=input,
27 |             weight=self.params["clf.linear.weight"],
28 |             bias=self.params["clf.linear.bias"],
29 |         )
30 | 
31 | 
32 | def extend_classification_head(model):
33 |     base_cls = model.classifier.__class__
34 |     base_cls_name = model.classifier.__class__.__name__
35 |     model.classifier.__class__ = type(
36 |         base_cls_name, (ClassificationHeadMixin, base_cls), {}
37 |     )
38 | 


--------------------------------------------------------------------------------
/inc_rel/meta/learner.py:
--------------------------------------------------------------------------------
  1 | import logging
  2 | import math
  3 | from collections import defaultdict
  4 | from typing import Dict, Union
  5 | 
  6 | import meta.utils as utils
  7 | import numpy as np
  8 | import torch
  9 | import torch.nn as nn
 10 | from sentence_transformers import util
 11 | 
 12 | 
 13 | class Learner(nn.Module):
 14 |     def __init__(
 15 |         self,
 16 |         base_model,
 17 |         exclude_classifier: bool,
 18 |         learning_rate: float,
 19 |         enable_warmup: bool,
 20 |         num_warmup_steps: int,
 21 |         total_training_steps: int,
 22 |         num_steps_per_batch: int,
 23 |         meta_module: str,
 24 |         loss_fn,
 25 |     ):
 26 |         super().__init__()
 27 |         self.base_model = base_model
 28 |         self.loss_fn = loss_fn()
 29 |         self.exclude_classifier = exclude_classifier
 30 |         self.learning_rate = learning_rate
 31 |         self.enable_warmup = enable_warmup
 32 |         self.meta_module = meta_module
 33 |         self.num_steps_per_batch = num_steps_per_batch
 34 |         if num_warmup_steps > 0:
 35 |             self.lr_scheduler = utils.InnerLRScheduler(
 36 |                 learning_rate, num_warmup_steps, total_training_steps
 37 |             )
 38 |         # TODO make adapter name variable
 39 | 
 40 |     def meta_module_iter(self, get_all: bool = False):
 41 | 
 42 |         if self.meta_module == "adapter":
 43 |             for layer in getattr(
 44 |                 self.base_model, "bert", self.base_model
 45 |             ).encoder.layer:
 46 |                 for adapter in layer.output.adapters.values():
 47 |                     yield adapter
 48 |         elif self.meta_module == "bias":
 49 |             for n, m in getattr(
 50 |                 self.base_model, "bert", self.base_model
 51 |             ).encoder.named_modules(
 52 |                 prefix="bert." if hasattr(self.base_model, "bert") else "" + "encoder"
 53 |             ):
 54 |                 if isinstance(m, nn.Linear):
 55 |                     assert m.bias.requires_grad, n
 56 |                     yield m
 57 | 
 58 |         if get_all or not self.exclude_classifier:
 59 |             assert self.base_model.classifier.weight.requires_grad
 60 |             assert self.base_model.classifier.bias.requires_grad
 61 |             yield self.base_model.classifier
 62 | 
 63 |     def get_params(self):
 64 |         params = {}
 65 |         for meta_module in self.meta_module_iter():
 66 |             params.update(meta_module.params)
 67 | 
 68 |         return params
 69 | 
 70 |     def set_params(self, params: Dict[str, torch.Tensor]):
 71 |         for meta_module in self.meta_module_iter():
 72 |             meta_module.set_params(params)
 73 | 
 74 |     def reset_params(self):
 75 |         for meta_module in self.meta_module_iter(get_all=False):
 76 |             meta_module.reset_params()
 77 | 
 78 |     def forward(
 79 |         self,
 80 |         input_ids,
 81 |         token_type_ids,
 82 |         attention_mask,
 83 |         params: Union[Dict[str, torch.Tensor], None] = None,
 84 |     ):
 85 |         if params:
 86 |             self.set_params(params)
 87 |         else:
 88 |             self.reset_params()
 89 | 
 90 |         out = self.base_model(
 91 |             input_ids=input_ids,
 92 |             token_type_ids=token_type_ids,
 93 |             attention_mask=attention_mask,
 94 |         )
 95 |         # out = out[0].view(-1)
 96 | 
 97 |         return out
 98 | 
 99 |     def train_step(
100 |         self,
101 |         input_ids,
102 |         token_type_ids,
103 |         attention_mask,
104 |         targets,
105 |         batch_size,
106 |         eval_final: bool,
107 |         architecture: str,
108 |         query_batch: Dict = None,
109 |     ):
110 |         results = defaultdict(list)
111 |         params = None
112 | 
113 |         self.reset_params()
114 |         for batch_i, batch_inputs, batch_targets in utils.batch_iter(
115 |             math.ceil(input_ids.size(0) / self.num_steps_per_batch),
116 |             input_ids,
117 |             token_type_ids,
118 |             attention_mask,
119 |             targets,
120 |             shuffle=True,
121 |         ):
122 |             if architecture == "ce":
123 |                 out = self(*batch_inputs, params=params)
124 |                 logits = out[0].view(-1)
125 | 
126 |             elif architecture == "bi":
127 |                 doc_out = self(*batch_inputs, params=params)
128 |                 doc_embeddings = doc_out[1]
129 |                 query_out = self(**query_batch, params=params)
130 |                 query_embedding = query_out[1]
131 | 
132 |                 logits = util.pytorch_cos_sim(query_embedding, doc_embeddings)[0]
133 |             else:
134 |                 raise ValueError(architecture)
135 | 
136 |             loss = self.loss_fn(logits, batch_targets)
137 | 
138 |             if not params:
139 |                 params = self.get_params()
140 |             self.base_model.zero_grad()
141 |             try:
142 |                 grads = torch.autograd.grad(
143 |                     loss, params.values(), create_graph=False, allow_unused=False
144 |                 )
145 | 
146 |             except Exception as err:
147 |                 print(params.keys())
148 |                 grads = torch.autograd.grad(
149 |                     loss, params.values(), create_graph=False, allow_unused=True
150 |                 )
151 |                 for (name, param), grad in zip(params.items(), grads):
152 |                     if grad is None or grad.sum() == 0:
153 |                         print("grad is zero for ", name)
154 |                 print(params.keys())
155 |                 raise err
156 | 
157 |             lr = (
158 |                 self.lr_scheduler.get_lr() if self.enable_warmup else self.learning_rate
159 |             )
160 |             updated_params = {}
161 |             for (name, param), grad in zip(params.items(), grads):
162 |                 updated_params[name] = param - lr * grad
163 | 
164 |             if self.enable_warmup:
165 |                 self.lr_scheduler.step()
166 | 
167 |             params = updated_params
168 | 
169 |             with torch.no_grad():
170 |                 results["loss"].append(loss.item())
171 |                 accuracy = utils.binary_accuracy_from_logits(batch_targets, logits)
172 |                 results["accuracy"].append(accuracy)
173 | 
174 |         if eval_final:
175 |             with torch.no_grad():
176 |                 logits = self(
177 |                     input_ids,
178 |                     token_type_ids,
179 |                     attention_mask,
180 |                     params=params,
181 |                 )
182 |                 if architecture == "ce":
183 |                     out = self(
184 |                         input_ids,
185 |                         token_type_ids,
186 |                         attention_mask,
187 |                         params=params,
188 |                     )
189 |                     logits = out[0].view(-1)
190 | 
191 |                 elif architecture == "bi":
192 |                     doc_out = self(
193 |                         input_ids, token_type_ids, attention_mask, params=params
194 |                     )
195 |                     doc_embeddings = doc_out[1]
196 |                     query_out = self(**query_batch, params=params)
197 |                     query_embedding = query_out[1]
198 | 
199 |                     logits = util.pytorch_cos_sim(query_embedding, doc_embeddings)[0]
200 |             loss = self.loss_fn(logits, targets)
201 |             accuracy = utils.binary_accuracy_from_logits(targets, logits)
202 |             results["final-loss"].append(loss.item())
203 |             results["final-accuracy"].append(accuracy)
204 | 
205 |         for k, v in results.items():
206 |             if isinstance(v, list):
207 |                 results[k] = np.mean(v)
208 | 
209 |         return params, results
210 | 


--------------------------------------------------------------------------------
/inc_rel/meta/meta_dataset.py:
--------------------------------------------------------------------------------
 1 | import random
 2 | from collections import defaultdict
 3 | from typing import Dict, List
 4 | 
 5 | from torch.utils.data import Dataset
 6 | 
 7 | 
 8 | def make_train_query_annotations(annotations):
 9 |     topic_ids = list(annotations.keys())
10 |     random.shuffle(topic_ids)
11 |     n = len(topic_ids) // 2
12 |     train_topic_ids, query_topic_ids = topic_ids[:n], topic_ids[-n:]
13 |     train_annotations = {
14 |         topic_id: annotations[topic_id] for topic_id in train_topic_ids
15 |     }
16 |     query_annotations = {
17 |         topic_id: annotations[topic_id] for topic_id in query_topic_ids
18 |     }
19 |     return train_annotations, query_annotations
20 | 
21 | 
22 | class MetaDataset(Dataset):
23 |     def __init__(self, annotations: Dict[str, List[Dict]]):
24 |         self.annotations = annotations
25 |         self.idx2topic_id = {i: topic_id for i, topic_id in enumerate(annotations)}
26 |         self.data = []
27 |         for topic_id, topic_annotations in self.annotations.items():
28 |             documents, doc_ids, labels = [], [], []
29 |             for annot in topic_annotations:
30 |                 query = annot["query"]
31 |                 documents.append(annot["document"])
32 |                 doc_ids.append(annot["doc_id"])
33 |                 labels.append(annot["label"])
34 | 
35 |             self.data.append(
36 |                 {
37 |                     "topic_id": [topic_id] * len(documents),
38 |                     "query": [query] * len(documents),
39 |                     "documents": documents,
40 |                     "doc_ids": doc_ids,
41 |                     "labels": labels,
42 |                 }
43 |             )
44 | 
45 |     def __getitem__(self, idx):
46 |         return self.data[idx]
47 | 
48 |     def __len__(self):
49 |         return len(self.data)
50 | 


--------------------------------------------------------------------------------
/inc_rel/meta/module_meta.py:
--------------------------------------------------------------------------------
 1 | from typing import Dict
 2 | 
 3 | import torch.nn as nn
 4 | 
 5 | 
 6 | class ModuleMetaMixin(nn.Module):
 7 |     def get_params(self):
 8 |         raise NotImplementedError()
 9 | 
10 |     def set_params(self, params):
11 |         self._params = params
12 | 
13 |     def reset_params(self):
14 |         self.params = self.get_params()
15 | 


--------------------------------------------------------------------------------
/inc_rel/meta/utils.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from sklearn.metrics import accuracy_score
 3 | 
 4 | 
 5 | def binary_accuracy_from_logits(y_true, logits):
 6 |     y_true = y_true.detach().cpu()
 7 |     logits = logits.detach().cpu()
 8 |     probs = torch.sigmoid(logits)
 9 |     accuracy = accuracy_score(y_true.view(-1).long(), probs.view(-1) >= 0.5)
10 |     return accuracy
11 | 
12 | 
13 | def batch_iter(
14 |     batch_size: int,
15 |     input_ids,
16 |     token_type_ids,
17 |     attention_mask,
18 |     targets,
19 |     shuffle: bool = False,
20 | ):
21 |     if shuffle:
22 |         p = torch.randperm(n=input_ids.size(0), device=input_ids.device)
23 |         input_ids = input_ids[p]
24 |         token_type_ids = token_type_ids[p]
25 |         attention_mask = attention_mask[p]
26 |         targets = targets[p]
27 | 
28 |     for batch_i in range(0, input_ids.size(0), batch_size):
29 |         batch_input_ids = input_ids[batch_i : batch_i + batch_size]
30 |         batch_token_type_ids = token_type_ids[batch_i : batch_i + batch_size]
31 |         batch_attention_mask = attention_mask[batch_i : batch_i + batch_size]
32 |         batch_targets = targets[batch_i : batch_i + batch_size]
33 | 
34 |         yield batch_i, (
35 |             batch_input_ids,
36 |             batch_token_type_ids,
37 |             batch_attention_mask,
38 |         ), batch_targets
39 | 
40 | 
41 | class InnerLRScheduler:
42 |     def __init__(self, lr, num_warmup_steps, num_training_steps):
43 |         self.lr = lr
44 |         self.current_step = 0
45 |         self.num_warmup_steps = num_warmup_steps
46 |         self.num_training_steps = num_training_steps
47 | 
48 |     def step(self):
49 |         self.current_step += 1
50 | 
51 |     def get_lr(self):
52 |         # https://huggingface.co/transformers/_modules/transformers/optimization.html#get_linear_schedule_with_warmup
53 |         if self.current_step < self.num_warmup_steps:
54 |             lmdb = float(self.current_step) / float(max(1, self.num_warmup_steps))
55 |         else:
56 |             lmdb = max(
57 |                 0.0,
58 |                 float(self.num_training_steps - self.current_step)
59 |                 / float(max(1, self.num_training_steps - self.num_warmup_steps)),
60 |             )
61 |         return lmdb * self.lr
62 | 


--------------------------------------------------------------------------------
/inc_rel/pre_train.py:
--------------------------------------------------------------------------------
  1 | import json
  2 | import os
  3 | import shutil
  4 | from collections import defaultdict
  5 | 
  6 | import simple_parsing
  7 | import tqdm
  8 | from args import PreTrain
  9 | from few_shot_trainer import FewShotTrainer
 10 | from pre_train_trainer import PreTrainTrainer
 11 | from reranking_evaluator import RerankingEvaluator
 12 | from utils import get_best_experiment
 13 | 
 14 | 
 15 | def main(args):
 16 | 
 17 |     ranking_evaluator = RerankingEvaluator(args.qrels)
 18 |     split2metric = defaultdict(list)
 19 | 
 20 |     for seed in args.seeds:
 21 |         print(f"---seed={seed}---")
 22 |         annotations = defaultdict(dict)
 23 |         for split in args.splits:
 24 |             with open(
 25 |                 os.path.join(
 26 |                     args.data_path, f"k{args.num_samples}", f"s{seed}", f"{split}.json"
 27 |                 )
 28 |             ) as fh:
 29 |                 annotations[split] = json.load(fh)
 30 | 
 31 |         # PRE-TRAIN model on training set
 32 |         results = []
 33 |         lr2best_metric = {}
 34 |         search_steps = args.epochs * len(args.learning_rates)
 35 |         hparam_results_file = args.hparam_results_file.format(seed=seed)
 36 |         with tqdm.tqdm(
 37 |             total=search_steps,
 38 |             ncols=100,
 39 |             desc="pre-train h-param search",
 40 |         ) as pbar:
 41 | 
 42 |             pre_train_trainer = PreTrainTrainer(
 43 |                 model_name=args.model,
 44 |                 ft_params=args.ft_params,
 45 |                 docs=args.bm25_docs,
 46 |                 initial_ranking=args.bm25_results,
 47 |                 ranking_evaluator=ranking_evaluator,
 48 |                 meta=args.pretrain_method == "meta",
 49 |                 pbar=pbar,
 50 |             )
 51 | 
 52 |             for learning_rate in args.learning_rates:
 53 |                 result, best_epoch, best_metric = pre_train_trainer.train(
 54 |                     train_annotations=annotations["train"],
 55 |                     eval_annotations=annotations["valid"],
 56 |                     epochs=args.epochs,
 57 |                     learning_rate=learning_rate,
 58 |                     exp_dir=os.path.join(
 59 |                         args.exp_path, f"s{seed}_lr-{learning_rate:.8f}"
 60 |                     ),
 61 |                     selection_metric=args.metric,
 62 |                 )
 63 |                 lr2best_metric[learning_rate] = best_metric
 64 |                 results.append({"result": result, "best_epoch": best_epoch})
 65 |                 # write out results after each experiment
 66 |                 with open(hparam_results_file, "w") as fh:
 67 |                     json.dump(results, fh, indent=4)
 68 | 
 69 |         best_lr = max(lr2best_metric, key=lr2best_metric.get)
 70 |         print(f"Pre-Train best LR={best_lr} {args.metric}={lr2best_metric[best_lr]:.4}")
 71 | 
 72 |         for learning_rate in args.learning_rates:
 73 |             exp_dir = os.path.join(args.exp_path, f"s{seed}_lr-{learning_rate:.8f}")
 74 |             if learning_rate == best_lr:
 75 |                 best_exp_dir = os.path.join(args.exp_path, f"best-model-s{seed}")
 76 |                 os.rename(exp_dir, best_exp_dir)
 77 |             else:
 78 |                 shutil.rmtree(exp_dir)
 79 | 
 80 |         # FINE-TUNE model on valid set
 81 |         search_steps = (
 82 |             args.epochs * len(args.learning_rates) * len(annotations["valid"])
 83 |         )
 84 |         with tqdm.tqdm(
 85 |             total=search_steps,
 86 |             ncols=100,
 87 |             desc="few-shot h-param search",
 88 |         ) as pbar:
 89 |             few_shot_trainer = FewShotTrainer(
 90 |                 model_name=args.model,
 91 |                 ft_params=args.ft_params,
 92 |                 docs=args.bm25_docs,
 93 |                 initial_ranking=args.bm25_results,
 94 |                 ranking_evaluator=ranking_evaluator,
 95 |                 pbar=pbar,
 96 |                 model_path=os.path.join(best_exp_dir, "model"),
 97 |             )
 98 | 
 99 |             results = []
100 |             for learning_rate in args.learning_rates:
101 |                 exp_results = few_shot_trainer.train(
102 |                     annotations["valid"],
103 |                     epochs=args.epochs,
104 |                     learning_rate=learning_rate,
105 |                 )
106 |                 results.extend(exp_results)
107 |                 # write out results after each experiment
108 |                 with open(hparam_results_file, "w") as fh:
109 |                     json.dump(results, fh, indent=4)
110 | 
111 |         # Get best model and eval on train/valid/test set
112 |         best_epochs, best_learning_rate = get_best_experiment(results, args.metric)
113 |         print(f"Few-Shot best Epoch={best_epoch:02d} best LR={best_learning_rate}")
114 | 
115 |         for split in args.splits:
116 |             few_shot_results = few_shot_trainer.train(
117 |                 topic2annotations=annotations[split],
118 |                 epochs=best_epochs,
119 |                 learning_rate=best_learning_rate,
120 |                 iter_epochs=False,
121 |                 update_pbar=False,
122 |             )
123 |             results_file = os.path.join(
124 |                 args.exp_path, f"k{args.num_samples}_s{seed}_{split}_results.json"
125 |             )
126 |             with open(results_file, "w") as fh:
127 |                 json.dump(few_shot_results, fh, indent=4)
128 | 
129 |             mean_metric = 0
130 |             for result in few_shot_results:
131 |                 mean_metric += (
132 |                     1
133 |                     / len(few_shot_results)
134 |                     * result["metrics"][list(result["metrics"].keys())[0]][args.metric]
135 |                 )
136 |             print(
137 |                 f"k={args.num_samples:02d} split={split:5s} "
138 |                 f"{args.metric}={mean_metric:.4f}"
139 |             )
140 |             split2metric[split].append(mean_metric)
141 | 
142 |     print("---MEAN---")
143 |     for split in args.splits:
144 |         print(
145 |             f"k={args.num_samples:02d} split={split:5s} "
146 |             f"{args.metric}={sum(split2metric[split]) / len(split2metric[split]):.4f}"
147 |         )
148 | 
149 | 
150 | if __name__ == "__main__":
151 |     args = simple_parsing.parse(PreTrain)
152 |     print(args)
153 |     main(args)
154 | 


--------------------------------------------------------------------------------
/inc_rel/pre_train_trainer.py:
--------------------------------------------------------------------------------
  1 | import json
  2 | import os
  3 | import random
  4 | from collections import defaultdict
  5 | from typing import Any, Dict, List
  6 | 
  7 | import eval
  8 | import numpy as np
  9 | import torch
 10 | import transformers
 11 | from meta.adapter_meta import extend_text_adapter
 12 | from meta.bias_meta import extend_linear
 13 | from meta.classification_head_meta import extend_classification_head
 14 | from meta.learner import Learner
 15 | from meta.meta_dataset import MetaDataset
 16 | from reranking_evaluator import RerankingEvaluator
 17 | from sentence_transformers import (
 18 |     CrossEncoder,
 19 |     InputExample,
 20 |     SentenceTransformer,
 21 |     losses,
 22 |     util,
 23 | )
 24 | from torch.nn.modules.loss import BCEWithLogitsLoss
 25 | from torch.utils.data.dataloader import DataLoader
 26 | from tqdm.auto import tqdm
 27 | 
 28 | # from transformers import AdapterConfig
 29 | 
 30 | 
 31 | class PreTrainTrainer:
 32 |     def __init__(
 33 |         self,
 34 |         model_name: str,
 35 |         ft_params: str,
 36 |         docs,
 37 |         initial_ranking,
 38 |         ranking_evaluator: RerankingEvaluator,
 39 |         meta,
 40 |         pbar,
 41 |     ):
 42 | 
 43 |         self.model_name = model_name
 44 |         self.ft_params = ft_params
 45 |         self.docs = docs
 46 |         self.initial_ranking = initial_ranking
 47 |         self.ranking_evaluator = ranking_evaluator
 48 |         self.meta = meta
 49 |         self.pbar = pbar
 50 | 
 51 |     def init_model(self):
 52 | 
 53 |         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 54 | 
 55 |         if self.model_name.startswith("cross-encoder"):
 56 |             model = CrossEncoder(self.model_name)
 57 |             model_ref = model.model.bert
 58 |             named_parameters = model.model.named_parameters
 59 |             parameters = model.model.parameters
 60 |             loss_fn = BCEWithLogitsLoss()
 61 | 
 62 |         elif self.model_name.startswith("sentence-transformer"):
 63 |             model = SentenceTransformer(self.model_name)
 64 |             model.to(device)
 65 |             model_ref = model[0].auto_model
 66 |             named_parameters = model.named_parameters
 67 |             parameters = model.parameters
 68 |             loss_fn = losses.CosineSimilarityLoss(model)
 69 |         else:
 70 |             ValueError(self.model_name)
 71 | 
 72 |         # init params to fine-tune
 73 | 
 74 |         if self.ft_params == "bias":
 75 |             for name, param in named_parameters():
 76 |                 if "bias" not in name:
 77 |                     param.requires_grad = False
 78 |             if self.meta:
 79 |                 extend_linear(
 80 |                     model_ref,
 81 |                     prefix="bert."
 82 |                     if isinstance(model, CrossEncoder)
 83 |                     else "" + "encoder",
 84 |                 )
 85 |                 if isinstance(model, CrossEncoder):
 86 |                     model.model.classifier.weight.requires_grad = True
 87 |                     extend_classification_head(model.model)
 88 | 
 89 |         elif self.ft_params == "head":
 90 |             if isinstance(model, SentenceTransformer):
 91 |                 raise RuntimeError("Head Tuning not compatiable with Bi-Encoder.")
 92 |             for name, param in named_parameters():
 93 |                 if "classifier" not in name:
 94 |                     param.requires_grad = False
 95 |             if self.meta:
 96 |                 extend_classification_head(model_ref)
 97 |         # elif self.ft_params == "adapter":
 98 |         #     torch.manual_seed(0)
 99 |         #     adapter_config = AdapterConfig.load("pfeiffer", reduction_factor=16)
100 |         #     adapter_name = "adapter"
101 |         #     model_ref.add_adapter(adapter_name, config=adapter_config)
102 |         #     model_ref.set_active_adapters(adapter_name)
103 |         #     model_ref.train_adapter(adapter_name)
104 |         #     if self.meta:
105 |         #         extend_text_adapter(model_ref)
106 |         #         if isinstance(model, CrossEncoder):
107 |         #             extend_classification_head(model.model)
108 | 
109 |         elif self.ft_params == "full":
110 |             pass
111 | 
112 |         else:
113 |             raise ValueError(self.ft_params)
114 | 
115 |         if self.model_name.startswith("cross-encoder"):
116 |             model.model.to(device)
117 |         elif self.model_name.startswith("sentence-transformer"):
118 |             model.to(device)
119 | 
120 |         return model, loss_fn, parameters, named_parameters
121 | 
122 |     def init_optimizer(self, named_parameters, learning_rate):
123 |         weight_decay = 0.01
124 |         no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
125 |         optimizer_grouped_parameters = [
126 |             {
127 |                 "params": [
128 |                     p
129 |                     for n, p in named_parameters()
130 |                     if not any(nd in n for nd in no_decay)
131 |                 ],
132 |                 "weight_decay": weight_decay,
133 |             },
134 |             {
135 |                 "params": [
136 |                     p for n, p in named_parameters() if any(nd in n for nd in no_decay)
137 |                 ],
138 |                 "weight_decay": 0.0,
139 |             },
140 |         ]
141 |         optimizer = transformers.AdamW(
142 |             optimizer_grouped_parameters, lr=learning_rate, no_deprecation_warning=True
143 |         )
144 |         return optimizer
145 | 
146 |     def init_dataloader(self, annotations, batch_size):
147 |         dataset = []
148 |         for annotation in annotations.values():
149 |             for a in annotation:
150 |                 dataset.append(
151 |                     InputExample(
152 |                         texts=[a["query"], a["document"]],
153 |                         label=float(a["label"] > 0),
154 |                     )
155 |                 )
156 |         g = torch.Generator()
157 |         g.manual_seed(0)
158 |         return DataLoader(dataset, batch_size=batch_size, shuffle=True, generator=g)
159 | 
160 |     @staticmethod
161 |     def batch_to_device(batch, device):
162 |         batch = {k: v.to(device) for k, v in batch.items()}
163 |         return batch
164 | 
165 |     def init_meta_dataloader(
166 |         self, annotations, batch_size, tokenizer, device, architecture: str
167 |     ):
168 |         def collate_fn(batch):
169 |             collated_batch = defaultdict(list)
170 |             for sample in batch:
171 |                 for k, v in sample.items():
172 |                     collated_batch[k].extend(v)
173 |             collated_batch = dict(collated_batch)
174 | 
175 |             if architecture == "ce":
176 |                 query_doc_sequences = tokenizer(
177 |                     collated_batch["query"],
178 |                     collated_batch["documents"],
179 |                     max_length=512,
180 |                     truncation="only_second",
181 |                     padding=True,
182 |                     return_tensors="pt",
183 |                 )
184 |                 query_doc_sequences = self.batch_to_device(query_doc_sequences, device)
185 |                 padded_sequences = (query_doc_sequences, None)
186 |             elif architecture == "bi":
187 |                 doc_sequences = tokenizer(
188 |                     collated_batch["documents"],
189 |                     max_length=512,
190 |                     truncation="only_first",
191 |                     padding=True,
192 |                     return_tensors="pt",
193 |                 )
194 |                 doc_sequences = self.batch_to_device(doc_sequences, device)
195 | 
196 |                 query_sequences = tokenizer(
197 |                     collated_batch["query"][0],
198 |                     max_length=512,
199 |                     truncation="only_first",
200 |                     padding=True,
201 |                     return_tensors="pt",
202 |                 )
203 |                 query_sequences = self.batch_to_device(query_sequences, device)
204 |                 padded_sequences = (doc_sequences, query_sequences)
205 |             else:
206 |                 raise ValueError(architecture)
207 | 
208 |             targets = torch.Tensor(collated_batch["labels"]).to(device)
209 |             targets = (targets > 0).float()
210 | 
211 |             return padded_sequences, targets
212 | 
213 |         g = torch.Generator()
214 |         g.manual_seed(0)
215 | 
216 |         return DataLoader(
217 |             MetaDataset(annotations),
218 |             batch_size=batch_size,
219 |             collate_fn=collate_fn,
220 |             shuffle=True,
221 |             generator=g,
222 |         )
223 | 
224 |     def train(self, *args, **kwargs):
225 |         if self.meta:
226 |             train_result = self.train_meta(*args, **kwargs)
227 |         else:
228 |             train_result = self.train_supvervised(*args, **kwargs)
229 |         return train_result
230 | 
231 |     def train_supvervised(
232 |         self,
233 |         train_annotations: Dict[str, List],
234 |         eval_annotations: Dict[str, List],
235 |         epochs: int,
236 |         learning_rate: float,
237 |         selection_metric: str,
238 |         exp_dir: str,
239 |         max_grad_norm: float = 1,
240 |         train_batch_size: int = 8,
241 |         show_pbar: bool = True,
242 |     ) -> List[Dict[str, Any]]:
243 | 
244 |         model, loss_fn, parameters, named_parameters = self.init_model()
245 | 
246 |         optimizer = self.init_optimizer(named_parameters, learning_rate)
247 | 
248 |         train_dataloader = self.init_dataloader(train_annotations, train_batch_size)
249 |         train_dataloader.collate_fn = model.smart_batching_collate
250 | 
251 |         epoch_results = []
252 |         best_epoch = 0
253 |         best_metric = 0
254 |         for epoch in range(epochs):
255 | 
256 |             for features, labels in train_dataloader:
257 |                 if isinstance(model, SentenceTransformer):
258 |                     loss_value = loss_fn(features, labels)
259 |                 elif isinstance(model, CrossEncoder):
260 |                     output = model.model(**features)
261 |                     loss_value = loss_fn(output.logits.view(-1), labels)
262 |                 loss_value.backward()
263 |                 torch.nn.utils.clip_grad_norm_(parameters(), max_grad_norm)
264 |                 optimizer.step()
265 |                 optimizer.zero_grad()
266 | 
267 |             epoch_result, best_epoch, best_metric = self.eval(
268 |                 model,
269 |                 eval_annotations,
270 |                 epoch,
271 |                 learning_rate,
272 |                 best_epoch,
273 |                 best_metric,
274 |                 selection_metric,
275 |                 exp_dir,
276 |             )
277 |             epoch_results.append(epoch_result)
278 |             self.pbar.update(1)
279 | 
280 |         return epoch_results, best_epoch, best_metric
281 | 
282 |     def train_meta(
283 |         self,
284 |         train_annotations: Dict[str, List],
285 |         eval_annotations: Dict[str, List],
286 |         epochs: int,
287 |         learning_rate: float,
288 |         selection_metric: str,
289 |         exp_dir: str,
290 |         max_grad_norm: float = 1,
291 |         train_batch_size: int = 8,
292 |         show_pbar: bool = True,
293 |     ):
294 |         model, loss_fn, parameters, named_parameters = self.init_model()
295 |         architecture = "ce" if isinstance(model, CrossEncoder) else "bi"
296 | 
297 |         optimizer = self.init_optimizer(named_parameters, learning_rate)
298 | 
299 |         learner = Learner(
300 |             base_model=model.model
301 |             if isinstance(model, CrossEncoder)
302 |             else model[0].auto_model,
303 |             exclude_classifier=isinstance(model, SentenceTransformer),
304 |             learning_rate=learning_rate,
305 |             enable_warmup=False,
306 |             num_warmup_steps=0,
307 |             total_training_steps=None,
308 |             num_steps_per_batch=1,
309 |             meta_module=self.ft_params,
310 |             loss_fn=torch.nn.BCEWithLogitsLoss
311 |             if architecture == "ce"
312 |             else torch.nn.MSELoss,
313 |         )
314 | 
315 |         train_annotations, query_annotations = self.make_train_query_annotations(
316 |             train_annotations
317 |         )
318 |         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
319 |         train_dataloader = self.init_meta_dataloader(
320 |             annotations=train_annotations,
321 |             batch_size=1,
322 |             tokenizer=model.tokenizer,
323 |             device=device,
324 |             architecture=architecture,
325 |         )
326 |         query_dataloader = self.init_meta_dataloader(
327 |             annotations=query_annotations,
328 |             batch_size=1,
329 |             tokenizer=model.tokenizer,
330 |             device=device,
331 |             architecture=architecture,
332 |         )
333 |         epoch_results = []
334 |         best_epoch = 0
335 |         best_metric = 0
336 |         for epoch in range(epochs):
337 |             for (train_batch, train_targets), (query_batch, query_targets) in zip(
338 |                 train_dataloader, query_dataloader
339 |             ):
340 |                 params, results = learner.train_step(
341 |                     **train_batch[0],
342 |                     targets=train_targets,
343 |                     batch_size=train_batch_size,
344 |                     eval_final=True,
345 |                     architecture=architecture,
346 |                     query_batch=train_batch[1],
347 |                 )
348 |                 if architecture == "ce":
349 |                     logits = learner(**query_batch[0], params=params,)[
350 |                         0
351 |                     ].view(-1)
352 |                     loss = loss_fn(logits, query_targets)
353 |                 elif architecture == "bi":
354 |                     doc_out = learner(**query_batch[0], params=params)
355 |                     doc_embeddings = doc_out[1]
356 |                     query_out = learner(**query_batch[1], params=params)
357 |                     query_embedding = query_out[1]
358 |                     logits = util.pytorch_cos_sim(query_embedding, doc_embeddings)[0]
359 |                     loss = torch.nn.MSELoss()(logits, query_targets)
360 | 
361 |                 optimizer.zero_grad()
362 |                 loss.backward()
363 |                 torch.nn.utils.clip_grad_norm_(parameters(), max_grad_norm)
364 |                 optimizer.step()
365 | 
366 |             epoch_result, best_epoch, best_metric = self.eval(
367 |                 model,
368 |                 eval_annotations,
369 |                 epoch,
370 |                 learning_rate,
371 |                 best_epoch,
372 |                 best_metric,
373 |                 selection_metric,
374 |                 exp_dir,
375 |             )
376 |             epoch_results.append(epoch_result)
377 |             self.pbar.update(1)
378 | 
379 |         return epoch_results, best_epoch, best_metric
380 | 
381 |     def eval(
382 |         self,
383 |         model,
384 |         annotations,
385 |         epoch,
386 |         learning_rate,
387 |         best_epoch,
388 |         best_metric,
389 |         selection_metric,
390 |         exp_dir,
391 |     ):
392 | 
393 |         queries = {
394 |             topic_id: annotation[0]["query"]  # query is the same in each annotaiton
395 |             for topic_id, annotation in annotations.items()
396 |         }
397 |         metrics, run = self.ranking_evaluator(
398 |             model=model,
399 |             queries=queries,
400 |             docs=self.docs,
401 |             initial_ranking=self.initial_ranking,
402 |         )
403 | 
404 |         if eval.accumulate_results(metrics)["mean"][selection_metric] > best_metric:
405 |             best_metric = eval.accumulate_results(metrics)["mean"][selection_metric]
406 |             best_epoch = epoch
407 |             os.makedirs(exp_dir, exist_ok=True)
408 |             model.save(os.path.join(exp_dir, "model"))
409 | 
410 |         result = {
411 |             "epoch": epoch,
412 |             "learning_rate": learning_rate,
413 |             "metrics": metrics,
414 |             "run": run,
415 |         }
416 |         if best_epoch == epoch:
417 |             with open(
418 |                 os.path.join(exp_dir, "valid_best_epoch_results.json"), "w"
419 |             ) as fh:
420 |                 json.dump(result, fh, indent=4)
421 | 
422 |         return result, best_epoch, best_metric
423 | 
424 |     @staticmethod
425 |     def make_train_query_annotations(annotations, seed=1):
426 |         topic_ids = list(annotations.keys())
427 |         random.seed(seed)
428 |         random.shuffle(topic_ids)
429 |         n = len(topic_ids) // 2
430 |         train_topic_ids, query_topic_ids = topic_ids[:n], topic_ids[-n:]
431 |         train_annotations = {
432 |             topic_id: annotations[topic_id] for topic_id in train_topic_ids
433 |         }
434 |         query_annotations = {
435 |             topic_id: annotations[topic_id] for topic_id in query_topic_ids
436 |         }
437 |         return train_annotations, query_annotations
438 | 


--------------------------------------------------------------------------------
/inc_rel/query_fine_tune.py:
--------------------------------------------------------------------------------
  1 | import json
  2 | import os
  3 | from collections import defaultdict
  4 | 
  5 | import simple_parsing
  6 | import tqdm
  7 | from args import FineTuneExperiment
  8 | from few_shot_trainer import FewShotTrainer
  9 | from reranking_evaluator import RerankingEvaluator
 10 | from utils import get_best_experiment
 11 | 
 12 | 
 13 | def main(args):
 14 | 
 15 |     ranking_evaluator = RerankingEvaluator(args.qrels)
 16 |     split2metric = defaultdict(list)
 17 |     for seed in args.seeds:
 18 |         print(f"---seed={seed}---")
 19 |         annotations = defaultdict(dict)
 20 |         for split in args.splits:
 21 |             with open(
 22 |                 os.path.join(
 23 |                     args.data_path, f"k{args.num_samples}", f"s{seed}", f"{split}.json"
 24 |                 )
 25 |             ) as fh:
 26 |                 annotations[split] = json.load(fh)
 27 | 
 28 |         search_steps = (
 29 |             args.epochs * len(args.learning_rates) * len(annotations["valid"])
 30 |         )
 31 | 
 32 |         with tqdm.tqdm(
 33 |             total=search_steps,
 34 |             ncols=100,
 35 |             desc=f"few-shot h-param search seed={seed}",
 36 |         ) as pbar:
 37 | 
 38 |             few_shot_trainer = FewShotTrainer(
 39 |                 model_name=args.model,
 40 |                 ft_params=args.ft_params,
 41 |                 docs=args.bm25_docs,
 42 |                 initial_ranking=args.bm25_results,
 43 |                 ranking_evaluator=ranking_evaluator,
 44 |                 pbar=pbar,
 45 |             )
 46 | 
 47 |             results = []
 48 |             for learning_rate in args.learning_rates:
 49 |                 exp_results = few_shot_trainer.train(
 50 |                     annotations["valid"], args.epochs, learning_rate
 51 |                 )
 52 |                 results.extend(exp_results)
 53 |                 # write out results after each experiment
 54 |                 with open(args.hparam_results_file.format(seed=seed), "w") as fh:
 55 |                     json.dump(results, fh, indent=4)
 56 | 
 57 |             best_epochs, best_learning_rate = get_best_experiment(
 58 |                 results, "ndcg_cut_20"
 59 |             )
 60 |             print(f"Best Epoch={best_epochs}. Best LR={best_learning_rate}")
 61 |             for split in args.splits:
 62 |                 few_shot_results = few_shot_trainer.train(
 63 |                     topic2annotations=annotations[split],
 64 |                     epochs=best_epochs,
 65 |                     learning_rate=best_learning_rate,
 66 |                     iter_epochs=False,
 67 |                     update_pbar=False,
 68 |                 )
 69 | 
 70 |                 results_file = os.path.join(
 71 |                     args.exp_path, f"k{args.num_samples}_s{seed}_{split}_results.json"
 72 |                 )
 73 |                 with open(results_file, "w") as fh:
 74 |                     json.dump(few_shot_results, fh, indent=4)
 75 | 
 76 |                 mean_metric = 0
 77 |                 for result in few_shot_results:
 78 |                     mean_metric += (
 79 |                         1
 80 |                         / len(few_shot_results)
 81 |                         * result["metrics"][list(result["metrics"].keys())[0]][
 82 |                             args.metric
 83 |                         ]
 84 |                     )
 85 |                 print(
 86 |                     f"k={args.num_samples:02d} split={split:5s} "
 87 |                     f"{args.metric}={mean_metric:.4f}"
 88 |                 )
 89 |                 split2metric[split].append(mean_metric)
 90 | 
 91 |     print("---MEAN---")
 92 |     for split in args.splits:
 93 |         print(
 94 |             f"k={args.num_samples:02d} split={split:5s} "
 95 |             f"{args.metric}={sum(split2metric[split]) / len(split2metric[split]):.4f}"
 96 |         )
 97 | 
 98 | 
 99 | if __name__ == "__main__":
100 |     args = simple_parsing.parse(FineTuneExperiment)
101 |     print(args)
102 |     main(args)
103 | 


--------------------------------------------------------------------------------
/inc_rel/rank_fusion.py:
--------------------------------------------------------------------------------
 1 | import json
 2 | import os
 3 | from collections import defaultdict
 4 | 
 5 | import simple_parsing
 6 | from args import RankFusion
 7 | from eval import accumulate_results
 8 | from reranking_evaluator import RerankingEvaluator
 9 | 
10 | 
11 | def rank_fusion(rank: int, c: int = 60):
12 |     return 1 / (c + rank)
13 | 
14 | 
15 | def main(args):
16 | 
17 |     # load result files
18 |     results = []
19 |     for file in args.result_files:
20 |         with open(file) as fh:
21 |             results.append(json.load(fh))
22 | 
23 |     # compute rank fusion scores
24 |     query_to_doc_to_scores = defaultdict(lambda: defaultdict(float))
25 |     for result in results:
26 |         for query_id, doc_to_score in result.items():
27 |             for rank, doc_id in enumerate(doc_to_score.keys(), start=1):
28 |                 query_to_doc_to_scores[query_id][doc_id] += rank_fusion(rank)
29 | 
30 |     ranking_evaluator = RerankingEvaluator(args.qrels)
31 |     result = ranking_evaluator.eval(query_to_doc_to_scores)
32 | 
33 |     split2metric = defaultdict(list)
34 |     for seed in args.seeds:
35 |         for split in args.splits:
36 | 
37 |             topic_ids = list(args.topic_ids_split_seed[split, seed].keys())
38 |             split_seed_eval_acc = accumulate_results(result, topic_ids=topic_ids)
39 | 
40 |             eval_acc_file = os.path.join(
41 |                 args.exp_path, f"k{args.num_samples}_s{seed}_{split}_eval_acc.json"
42 |             )
43 |             with open(eval_acc_file, "w") as fh:
44 |                 json.dump(split_seed_eval_acc, fh, indent=4)
45 | 
46 |             split2metric[split].append(split_seed_eval_acc["mean"][args.metric])
47 | 
48 |             print(f"---seed={seed}---")
49 |             print(
50 |                 f"k={args.num_samples:02d} split={split:5s} seed={seed:02d} "
51 |                 f"{args.metric}={split_seed_eval_acc['mean'][args.metric]:.4f}"
52 |             )
53 | 
54 |     print("---MEAN---")
55 |     for split in args.splits:
56 |         print(
57 |             f"k={args.num_samples:02d} split={split:5s} "
58 |             f"{args.metric}={sum(split2metric[split]) / len(split2metric[split]):.4f}"
59 |         )
60 | 
61 | 
62 | if __name__ == "__main__":
63 |     args = simple_parsing.parse(RankFusion)
64 |     print(args)
65 |     main(args)
66 | 


--------------------------------------------------------------------------------
/inc_rel/reranking_evaluator.py:
--------------------------------------------------------------------------------
  1 | from collections import defaultdict
  2 | from contextlib import contextmanager
  3 | from typing import Dict, List, Tuple, Union
  4 | 
  5 | import numpy as np
  6 | import torch
  7 | from pytrec_eval import RelevanceEvaluator
  8 | from sentence_transformers import util
  9 | from sentence_transformers import util as st_util
 10 | from sentence_transformers.cross_encoder.CrossEncoder import CrossEncoder
 11 | from sentence_transformers.SentenceTransformer import SentenceTransformer
 12 | from tqdm.auto import tqdm
 13 | 
 14 | 
 15 | # https://discuss.pytorch.org/t/opinion-eval-should-be-a-context-manager/18998/3
 16 | @contextmanager
 17 | def evaluating(model):
 18 |     """Temporarily switch to evaluation mode."""
 19 |     istrain = model.training
 20 |     try:
 21 |         model.eval()
 22 |         yield model
 23 |     finally:
 24 |         if istrain:
 25 |             model.train()
 26 | 
 27 | 
 28 | class RerankingEvaluator:
 29 |     def __init__(self, qrels: Dict[str, Dict[str, int]]):
 30 | 
 31 |         self.qrels = qrels
 32 |         at_k = [1, 5, 10, 20, 50, 100, 1000]
 33 |         map = "map_cut." + ",".join([str(k) for k in at_k])
 34 |         ndcg = "ndcg_cut." + ",".join([str(k) for k in at_k])
 35 |         recall = "recall." + ",".join([str(k) for k in at_k])
 36 |         precision = "P." + ",".join([str(k) for k in at_k])
 37 |         rprecision = "Rprec"
 38 |         self.measures = {map, ndcg, recall, precision, rprecision}
 39 | 
 40 |     def __call__(
 41 |         self,
 42 |         model: Union[CrossEncoder, SentenceTransformer],
 43 |         queries: Dict[str, str],
 44 |         docs: Dict[str, str],
 45 |         initial_ranking: Dict[str, Dict[str, float]],
 46 |         model_ctx: SentenceTransformer = None,
 47 |         batch_size: int = 32,
 48 |         tokenizer=None,
 49 |         show_progress_bar=False,
 50 |         scoring_fn="cos",
 51 |     ):
 52 | 
 53 |         if tokenizer is None:
 54 |             tokenizer = model.tokenizer
 55 | 
 56 |         if scoring_fn == "cos":
 57 |             scoring_fn = st_util.cos_sim
 58 |         elif scoring_fn == "dot":
 59 |             scoring_fn = st_util.dot_score
 60 | 
 61 |         run = defaultdict(lambda: defaultdict(dict))
 62 |         with torch.no_grad(), evaluating(
 63 |             model if isinstance(model, SentenceTransformer) else model.model
 64 |         ):
 65 |             for topic_id, query in tqdm(
 66 |                 queries.items(),
 67 |                 desc="eval queries",
 68 |                 ncols=100,
 69 |                 disable=not show_progress_bar,
 70 |             ):
 71 | 
 72 |                 doc_ids = initial_ranking[topic_id].keys()
 73 |                 query_docs = [docs[doc_id] for doc_id in doc_ids]
 74 |                 if isinstance(model, CrossEncoder):
 75 |                     sentences = [[query, doc] for doc in query_docs]
 76 |                     scores = model.predict(
 77 |                         sentences,
 78 |                         batch_size=batch_size,
 79 |                         show_progress_bar=show_progress_bar,
 80 |                     )
 81 |                 elif isinstance(model, SentenceTransformer):
 82 |                     query_embedding = model.encode(
 83 |                         query,
 84 |                         batch_size=batch_size,
 85 |                         show_progress_bar=show_progress_bar,
 86 |                     )
 87 |                     if model_ctx:
 88 |                         doc_model = model_ctx
 89 |                     else:
 90 |                         doc_model = model
 91 |                     doc_embeddings = doc_model.encode(
 92 |                         query_docs,
 93 |                         batch_size=batch_size,
 94 |                         show_progress_bar=show_progress_bar,
 95 |                     )
 96 |                     scores = scoring_fn(query_embedding, doc_embeddings)[0]
 97 |                 # sort results
 98 |                 run[topic_id] = {
 99 |                     doc_id: float(score)
100 |                     for doc_id, score in sorted(
101 |                         zip(doc_ids, scores),
102 |                         key=lambda doc_score: doc_score[1],
103 |                         reverse=True,
104 |                     )
105 |                 }
106 | 
107 |         # RelevanceEvaluator seems to swallow some metrics when re-used
108 |         # re-initializing fixes this.
109 |         results = self.eval(run)
110 | 
111 |         return results, run
112 | 
113 |     def eval(self, run):
114 |         return RelevanceEvaluator(self.qrels, self.measures).evaluate(run)
115 | 


--------------------------------------------------------------------------------
/inc_rel/settings.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | from typing import List, Union
  3 | 
  4 | from pydantic import BaseSettings, Field, SecretStr
  5 | 
  6 | 
  7 | class Settings(BaseSettings):
  8 |     class Config:
  9 |         env_file = ".env"
 10 |         env_prefix = "INC_REL_"
 11 | 
 12 |     ir_datasets_home: str = Field("~/.ir_datasets", env="IR_DATASETS_HOME")
 13 | 
 14 |     base_path: str
 15 |     raw_dir: str
 16 |     data_dir: str
 17 | 
 18 |     @property
 19 |     def raw_path(self) -> str:
 20 |         return os.path.join(self.base_path, self.raw_dir)
 21 | 
 22 |     @property
 23 |     def data_path(self) -> str:
 24 |         return os.path.join(self.base_path, self.data_dir)
 25 | 
 26 | 
 27 | class DatasetSettings(BaseSettings):
 28 | 
 29 |     split_seeds: List[int] = Field(
 30 |         [0, 1, 2], description="Random seeds for splitting the dataset."
 31 |     )
 32 |     split_sizes: List[float] = Field(
 33 |         [0.6, 0.2, 0.2], description="Train, dev, test split sizes."
 34 |     )
 35 |     num_samples: List[int] = Field(
 36 |         [2, 4, 8], description="Amount of relevance feedback per topic (k)."
 37 |     )
 38 |     bm25_size: int = Field(
 39 |         1000,
 40 |         description="Number of documents to retrieve in the first stage with BM25.",
 41 |     )
 42 |     remove_duplicates: bool = Field(False, description="Remove duplicate documents.")
 43 | 
 44 |     @property
 45 |     def corpus_path(self) -> str:
 46 |         # datasets managed by ir_datasets do not have a corpus_path
 47 |         return "not implemented"
 48 | 
 49 |     @property
 50 |     def topics_path(self) -> str:
 51 |         return "not implemented"
 52 | 
 53 |     @property
 54 |     def qrels_path(self) -> str:
 55 |         return "not implemented"
 56 | 
 57 |     @property
 58 |     def enrich_bm25_path(self) -> Union[str, None]:
 59 |         return None
 60 | 
 61 | 
 62 | class TRECCovidDatasetSettings(DatasetSettings):
 63 |     class Config:
 64 |         env_file = ".env"
 65 |         env_prefix = "INC_REL_TREC_COVID_"
 66 | 
 67 |     name: str = "trec-covid"
 68 |     raw_path: str
 69 |     data_path: str
 70 |     corpus_url: str
 71 |     topics_url: str
 72 |     qrels_url: str
 73 | 
 74 |     remove_duplicates = True
 75 | 
 76 |     @property
 77 |     def corpus_path(self) -> str:
 78 |         return os.path.join(self.raw_path, "2020-07-16", "metadata.csv")
 79 | 
 80 |     @property
 81 |     def topics_path(self) -> str:
 82 |         return os.path.join(self.raw_path, "topics-rnd5.xml")
 83 | 
 84 |     @property
 85 |     def qrels_path(self) -> str:
 86 |         return os.path.join(self.raw_path, "qrels-covid_d5_j0.5-5.txt")
 87 | 
 88 | 
 89 | class WebisTouche2020DatasetSettings(DatasetSettings):
 90 |     class Config:
 91 |         env_file = ".env"
 92 |         env_prefix = "INC_REL_TOUCHE_"
 93 | 
 94 |     name: str = "touche"
 95 |     ir_datasets_name: str
 96 |     data_path: str
 97 | 
 98 |     @property
 99 |     def enrich_bm25_path(self) -> str:
100 |         return os.path.join(
101 |             os.path.join(self.data_path, "1000", "full_bm25_results.json")
102 |         )
103 | 
104 | 
105 | class TRECRobustDatasetSettings(DatasetSettings):
106 |     class Config:
107 |         env_file = ".env"
108 |         env_prefix = "INC_REL_TREC_ROBUST_"
109 | 
110 |     name: str = "robust04"
111 |     raw_path: str
112 |     data_path: str
113 |     username: str
114 |     password: SecretStr
115 |     corpus_disk4_url: str
116 |     corpus_disk5_url: str
117 | 
118 | 
119 | class TRECNewsDatasetSettings(DatasetSettings):
120 |     class Config:
121 |         env_file = ".env"
122 |         env_prefix = "INC_REL_TREC_NEWS_"
123 | 
124 |     name: str = "trec-news"
125 |     raw_path: str
126 |     data_path: str
127 |     username: str
128 |     password: SecretStr
129 |     corpus_url: str
130 |     topics_url: str
131 |     qrels_url: str
132 | 
133 |     @property
134 |     def corpus_path(self) -> str:
135 |         return os.path.join(
136 |             self.raw_path,
137 |             "WashingtonPost.v2",
138 |             "data",
139 |             "TREC_Washington_Post_collection.v2.jl",
140 |         )
141 | 
142 |     @property
143 |     def topics_path(self) -> str:
144 |         return os.path.join(self.raw_path, "newsir19-background-linking-topics.xml")
145 | 
146 |     @property
147 |     def qrels_path(self) -> str:
148 |         return os.path.join(self.raw_path, "newsir19-qrels-background.txt")
149 | 
150 | 
151 | dataset_settings_cls = {
152 |     "trec-covid": TRECCovidDatasetSettings,
153 |     "robust04": TRECRobustDatasetSettings,
154 |     "trec-news": TRECNewsDatasetSettings,
155 |     "touche": WebisTouche2020DatasetSettings,
156 | }
157 | 


--------------------------------------------------------------------------------
/inc_rel/utils.py:
--------------------------------------------------------------------------------
  1 | import json
  2 | import os
  3 | import time
  4 | from html.parser import HTMLParser
  5 | from io import StringIO
  6 | from pathlib import Path
  7 | from typing import Callable, Dict, List, Tuple
  8 | 
  9 | import Levenshtein as lev
 10 | import numpy as np
 11 | import pandas as pd
 12 | from sklearn.model_selection import train_test_split
 13 | 
 14 | 
 15 | def get_best_experiment(results, selection_metric: str = "ndcg_cut_20"):
 16 |     for i, r in enumerate(results):
 17 |         for metric, score in r["metrics"][list(r["metrics"].keys())[0]].items():
 18 |             results[i][metric] = score
 19 |         r.pop("metrics")
 20 |         r.pop("run")
 21 |     df = pd.DataFrame(results)
 22 |     epoch, learning_rate = (
 23 |         df.groupby(["epoch", "learning_rate"])
 24 |         .mean(numeric_only=True)[selection_metric]
 25 |         .idxmax()
 26 |     )
 27 |     return epoch, learning_rate
 28 | 
 29 | 
 30 | def create_split(
 31 |     topic_ids: List[str], seed: int, split_sizes: List[float]
 32 | ) -> Tuple[List[str], List[str], List[str]]:
 33 | 
 34 |     if not sum(split_sizes) == 1:
 35 |         raise ValueError(
 36 |             f"Expected `split_sizes` to sum up to 1, but got {sum(split_sizes)=}"
 37 |         )
 38 | 
 39 |     split_sizes = list(map(lambda s: int(s * len(topic_ids)), split_sizes))
 40 |     while sum(split_sizes) < len(topic_ids):
 41 |         split_sizes[0] += 1
 42 | 
 43 |     train_topic_ids, remaining_topic_ids = train_test_split(
 44 |         topic_ids, train_size=split_sizes[0], random_state=seed
 45 |     )
 46 |     valid_topic_ids, test_topic_ids = train_test_split(
 47 |         remaining_topic_ids, train_size=split_sizes[1], random_state=seed
 48 |     )
 49 | 
 50 |     return train_topic_ids, valid_topic_ids, test_topic_ids
 51 | 
 52 | 
 53 | def get_similarity(docs: List[str], similarity_fn: Callable) -> np.ndarray:
 54 | 
 55 |     t1 = time.time()
 56 |     ratios = np.zeros((len(docs), (len(docs))))
 57 |     for i, i_doc in enumerate(docs):
 58 |         for j, j_doc in enumerate(docs[i:], start=i):
 59 |             if i == j:
 60 |                 continue
 61 |             ratios[i, j] = similarity_fn(i_doc.lower(), j_doc.lower())
 62 |     t2 = time.time()
 63 |     return ratios
 64 | 
 65 | 
 66 | def sorted_annotated_retrieved_docs(
 67 |     corpus: Dict[str, str],
 68 |     qrels: Dict[str, Dict[str, float]],
 69 |     bm25_results: Dict[str, Dict[str, float]],
 70 |     min_relevant_relevance: int,
 71 |     remove_duplicates: bool = False,
 72 |     enrich_negatives_from_bm25: Path = None,
 73 |     num_non_relevant: int = None,
 74 | ) -> Dict[str, Dict[str, List[str]]]:
 75 | 
 76 |     if enrich_negatives_from_bm25:
 77 |         with open(enrich_negatives_from_bm25) as fh:
 78 |             full_bm25 = json.load(fh)
 79 |         bm25_negatives = {}
 80 |         for topic_id, doc2score in full_bm25.items():
 81 |             bm25_negatives[topic_id] = list(doc2score.keys())[100:]
 82 | 
 83 |     candidates = {}
 84 | 
 85 |     for topic_id, annotated_doc_relevance in qrels.items():
 86 |         # annotated docs that were retrieved with bm25
 87 |         doc_ids = set(annotated_doc_relevance.keys()) & set(
 88 |             bm25_results[topic_id].keys()
 89 |         )
 90 |         if remove_duplicates:
 91 |             docs = [corpus[doc_id] for doc_id in doc_ids]
 92 | 
 93 |             similarity_ratio = get_similarity(docs, lambda x, y: int(x == y))
 94 |             duplicates = np.argwhere(similarity_ratio > 0.9)
 95 |             docs_to_remove = sorted(set(duplicates[:, 1].tolist()), reverse=True)
 96 |             for index in docs_to_remove:
 97 |                 del docs[index]
 98 | 
 99 |         # get all docs that were retrieved by bm25 and annotated with at least
100 |         # `min_relveant_relevance` or 0 for non relevant
101 |         relevant_docs = sorted(
102 |             filter(
103 |                 lambda doc_id: qrels[topic_id][doc_id] >= min_relevant_relevance,
104 |                 doc_ids,
105 |             ),
106 |             key=lambda doc_id: bm25_results[topic_id][doc_id],
107 |             reverse=True,
108 |         )
109 |         non_relevant_docs = sorted(
110 |             filter(
111 |                 lambda doc_id: qrels[topic_id][doc_id] <= 0,
112 |                 doc_ids,
113 |             ),
114 |             key=lambda doc_id: bm25_results[topic_id][doc_id],
115 |             reverse=True,
116 |         )
117 |         if len(non_relevant_docs) < num_non_relevant and enrich_negatives_from_bm25:
118 |             num_to_append = num_non_relevant - len(non_relevant_docs)
119 |             print(f"Appended {num_to_append} bm25 negatives to {topic_id=}.")
120 |             non_relevant_docs.extend(bm25_negatives[topic_id[:num_to_append]])
121 | 
122 |         candidates[topic_id] = {
123 |             "relevant": relevant_docs,
124 |             "non_relevant": non_relevant_docs,
125 |         }
126 | 
127 |     return candidates
128 | 
129 | 
130 | class MLStripper(HTMLParser):
131 |     def __init__(self):
132 |         super().__init__()
133 |         self.reset()
134 |         self.strict = False
135 |         self.convert_charrefs = True
136 |         self.text = StringIO()
137 | 
138 |     def handle_data(self, d):
139 |         self.text.write(d)
140 | 
141 |     def get_data(self):
142 |         return self.text.getvalue()
143 | 
144 | 
145 | def strip_tags(html):
146 |     s = MLStripper()
147 |     s.feed(html)
148 |     return s.get_data()
149 | 


--------------------------------------------------------------------------------
/inc_rel/zero_shot.py:
--------------------------------------------------------------------------------
 1 | import json
 2 | import os
 3 | from collections import defaultdict
 4 | 
 5 | import simple_parsing
 6 | from args import ZeroShot
 7 | from eval import accumulate_results
 8 | from reranking_evaluator import RerankingEvaluator
 9 | 
10 | 
11 | def main(args):
12 | 
13 |     model = args.model_class(args.model)
14 |     if args.model_ctx is not None:
15 |         model_ctx = args.model_class(args.model_ctx)
16 |     else:
17 |         model_ctx = None
18 | 
19 |     eval_results, run = RerankingEvaluator(args.qrels)(
20 |         model=model,
21 |         model_ctx=model_ctx,
22 |         queries=args.topics,
23 |         docs=args.bm25_docs,
24 |         initial_ranking=args.bm25_results,
25 |         batch_size=128,
26 |         show_progress_bar=False,
27 |         scoring_fn=args.scoring_fn,
28 |     )
29 | 
30 |     eval_file = os.path.join(args.exp_path, f"k{args.num_samples}_eval.json")
31 |     with open(eval_file, "w") as fh:
32 |         json.dump(eval_results, fh, indent=4)
33 | 
34 |     eval_acc_file = os.path.join(args.exp_path, f"k{args.num_samples}_eval_acc.json")
35 |     with open(
36 |         eval_acc_file,
37 |         "w",
38 |     ) as fh:
39 |         json.dump(accumulate_results(eval_results), fh, indent=4)
40 | 
41 |     results_file = os.path.join(args.exp_path, f"k{args.num_samples}_results.json")
42 |     with open(results_file, "w") as fh:
43 |         json.dump(run, fh, indent=4)
44 | 
45 |     split2metric = defaultdict(list)
46 |     for seed in args.seeds:
47 |         for split in args.splits:
48 |             topic_ids = list(args.topic_ids_split_seed[split, seed].keys())
49 |             split_seed_eval_acc = accumulate_results(eval_results, topic_ids=topic_ids)
50 | 
51 |             split_seed_eval_acc_file = os.path.join(
52 |                 args.exp_path, f"k{args.num_samples}_s{seed}_{split}_eval_acc.json"
53 |             )
54 |             with open(split_seed_eval_acc_file, "w") as fh:
55 |                 json.dump(split_seed_eval_acc, fh, indent=4)
56 | 
57 |             split2metric[split].append(split_seed_eval_acc["mean"][args.metric])
58 |             print(
59 |                 f"k={args.num_samples:02d} split={split:5s} seed={seed:02d} "
60 |                 f"{args.metric}={split_seed_eval_acc['mean'][args.metric]:.4f}"
61 |             )
62 | 
63 |     print("---MEAN---")
64 |     for split in args.splits:
65 |         print(
66 |             f"k={args.num_samples:02d} split={split:5s} "
67 |             f"{args.metric}={sum(split2metric[split]) / len(split2metric[split]):.4f}"
68 |         )
69 | 
70 | 
71 | if __name__ == "__main__":
72 |     args = simple_parsing.parse(ZeroShot)
73 |     print(args)
74 |     main(args)
75 | 


--------------------------------------------------------------------------------
/requirements.dev.txt:
--------------------------------------------------------------------------------
1 | black
2 | isort
3 | pre-commit
4 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | pydantic[dotenv]
 2 | requests
 3 | tqdm
 4 | simple-parsing
 5 | ir_datasets==0.4.0
 6 | Levenshtein==0.12.0
 7 | scikit-learn==0.24.2
 8 | elasticsearch==7.13.1
 9 | pytrec_eval==0.5
10 | sentence-transformers==2.1.0
11 | torch==1.11
12 | pandas==1.5.1
13 | 


--------------------------------------------------------------------------------