cloned most of ctmatch into this spaces repo... it will have to handle the data too, we'll see. using ctmatch requirements.txt

app.py

@@ -0,0 +1,27 @@
+
+from ctmatch.match import CTMatch, PipeConfig
+import gradio as gr
+
+
+pipe_config = PipeConfig(
+    classifier_model_checkpoint='semaj83/scibert_finetuned_pruned_ctmatch',
+    ir_setup=True,
+    filters=["svm", "classifier"],
+)
+
+CTM = CTMatch(pipe_config)
+
+
+def ctmatch_web_api(topic_query: str) -> str:
+    return '\n\n'.join([f"{nid}: {txt}" for nid, txt in CTM.match_pipeline(topic_query, top_k=5)])
+
+
+if __name__ == "__main__":
+  
+    with gr.Blocks(css=".gradio-container {background-color: #00CED1}") as demo:
+        name = gr.Textbox(lines=5, label="patient description", placeholder="Patient is a 45-year-old man with a history of anaplastic astrocytoma...")
+        output = gr.Textbox(lines=10, label="matching trials")
+        greet_btn = gr.Button("match")
+        greet_btn.click(fn=ctmatch_web_api, inputs=name, outputs=output, api_name="match")
+
+    demo.queue().launch(share=True, debug=True)

ctmatch/ct_data_paths.py

@@ -0,0 +1,92 @@
+
+from typing import List, Tuple
+
+TREC_REL_PATH =  "/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_21_judgments.txt"
+KZ_REL_PATH =  "/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/qrels-clinical_trials.txt"
+
+TREC_RELLED_TOPIC_PATH = "/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec21_topics.jsonl" 
+KZ_RELLED_TOPIC_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/processed_kz_data/processed_kz_topics.jsonl' 
+
+KZ_DOC_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/clinicaltrials.gov-16_dec_2015.zip'
+KZ_PROCESSED_DOC_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/processed_kz_data/processed_kz_docs.jsonl'
+
+TREC_ML_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_data.jsonl'
+KZ_ML_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/kz_data.jsonl'
+
+
+def get_data_tuples(trec_or_kz: str = 'trec') -> List[Tuple[str, str]]:
+	if trec_or_kz == 'trec':
+		return get_trec_doc_data_tuples(), get_trec_topic_data_tuples()
+	return get_kz_doc_data_tuples(), get_kz_topic_data_tuples()
+
+
+
+# --------------------------------------------------------------------------------------------------------------- #
+# data from TREC clinical track 2021 & 2022
+# --------------------------------------------------------------------------------------------------------------- #
+
+
+def get_trec_doc_data_tuples() -> List[Tuple[str]]:
+	trec22_pt1_docs = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_docs_21/ClinicalTrials.2021-04-27.part1.zip'
+	trec_pt1_target = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec22_docs_part1.jsonl'
+
+	trec22_pt2_docs = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_docs_21/ClinicalTrials.2021-04-27.part2.zip'
+	trec_pt2_target = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec22_docs_part2.jsonl'
+
+	trec22_pt3_docs = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_docs_21/ClinicalTrials.2021-04-27.part3.zip'
+	trec_pt3_target = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec22_docs_part3.jsonl'
+
+	trec22_pt4_docs = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_docs_21/ClinicalTrials.2021-04-27.part4.zip'
+	trec_pt4_target = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec22_docs_part4.jsonl'
+
+	trec22_pt5_docs = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_docs_21/ClinicalTrials.2021-04-27.part5.zip'
+	trec_pt5_target = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec22_docs_part5.jsonl'
+
+	trec_doc_data_tuples = [
+		(trec22_pt1_docs, trec_pt1_target), 
+		(trec22_pt2_docs, trec_pt2_target), 
+		(trec22_pt3_docs, trec_pt3_target),
+		(trec22_pt4_docs, trec_pt4_target),
+		(trec22_pt5_docs, trec_pt5_target)
+	]
+
+	return trec_doc_data_tuples
+
+
+def get_trec_topic_data_tuples() -> List[Tuple[str]]:
+	trec21_topic_path = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_21_topics.xml'
+	trec21_topic_target = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec21_topics.jsonl'
+	trec22_topic_path = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/trec_22_topics.xml'
+	trec22_topic_target = '/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec22_topics.jsonl'
+
+	trec_topic_data_tuples = [
+		(trec21_topic_path, trec21_topic_target), 
+		(trec22_topic_path, trec22_topic_target)
+	]
+	return trec_topic_data_tuples
+
+
+
+
+# --------------------------------------------------------------------------------------------------------------- #
+# data from Koontz, et al. (2016)
+# --------------------------------------------------------------------------------------------------------------- #
+def get_kz_doc_data_tuples() -> List[Tuple[str]]:
+	# kz_doc_data_tuples = []
+	# for i in range(1, 18):
+	# 	kz_doc_path = f'/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/kz_doc_splits/kz_doc_split{i}.zip'
+	# 	kz_doc_target = f'/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/processed_kz_data/processed_kz_doc_split{i}.jsonl'
+	# 	kz_doc_data_tuples.append((kz_doc_path, kz_doc_target))
+	kz_docs = KZ_DOC_PATH
+	kz_docs_target = KZ_PROCESSED_DOC_PATH
+	return [(kz_docs, kz_docs_target)]
+	
+	#return kz_doc_data_tuples
+
+def get_kz_topic_data_tuples() -> List[Tuple[str]]:
+	kz_topic_desc_path = '/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/topics-2014_2015-description.topics'
+	kz_topic_target = '/Users/jameskelly/Documents/cp/ctmatch/data/kz_data/processed_kz_data/processed_kz_topics.jsonl'
+	kz_topic_data_tuples = [
+		(kz_topic_desc_path, kz_topic_target)
+	]
+	return kz_topic_data_tuples

ctmatch/ctmatch_prep.py

@@ -0,0 +1,351 @@
+
+from typing import Dict, List, NamedTuple, Optional, Set, Tuple, Union
+import ct_data_paths as ctpaths
+import numpy as np
+import random
+import json
+
+from proc import CTConfig, CTProc, CTDocument, CTTopic
+from scripts.vis_scripts import analyze_test_rels
+from ctproc_ctmatch_utils import get_processed_data, truncate
+import ctproc_eda as eda
+
+LLM_END_PROMPT: str = "Revelance score (0, 1, or 2) : [CLS] "
+
+class DataConfig(NamedTuple):
+	save_path: str
+	trec_or_kz: str = 'trec'
+	filtered_topic_keys: Set[str] = {'id', 'text_sents', 'age', 'gender'}
+	filtered_doc_keys: Set[str] = {'id', 'elig_min_age', 'elig_max_age', 'elig_gender', 'condition', 'elig_crit'}
+	max_topic_len: Optional[int] = None
+	max_inc_len: Optional[int] = None
+	max_exc_len: Optional[int] = None
+	prepend_elig_age: bool = True
+	prepend_elig_gender: bool = True
+	include_only: bool = False
+	downsample_zeros_n: Optional[int] = None
+	sep: str = '[SEP]'
+	llm_prep: bool = False
+	first_n_only: Optional[int] = None
+	convert_snli: bool = False
+	infer_category_model: Optional[str] = None
+
+
+
+
+def proc_docs_and_topics(trec_or_kz: str = 'trec') -> Tuple[Dict[str, Dict[str, str]], Dict[str, Dict[str, str]]]:
+
+	doc_tuples, topic_tuples = ctpaths.get_data_tuples(trec_or_kz)
+
+	id2topic = dict()
+	for topic_source, topic_target in topic_tuples:
+		id2topic.update(proc_topics(topic_source, topic_target, trec_or_kz=trec_or_kz))
+		print(f"processed {trec_or_kz} topic source: {topic_source}, and wrote to {topic_target}")
+	
+	id2doc = dict()
+	for doc_source, doc_target in doc_tuples:
+		id2doc.update(proc_docs(doc_source, doc_target))
+		print(f"processed {trec_or_kz} doc source: {doc_source}, and wrote to {doc_target}")
+
+
+	return id2topic, id2doc
+
+
+
+
+
+def proc_docs(doc_path: str, output_path: str) -> Dict[str, CTDocument]:
+
+	ct_config = CTConfig(
+		data_path=doc_path, 
+		write_file=output_path,
+    	nlp=True
+	)
+
+	cp = CTProc(ct_config)
+	id2doc = {res.id : res for res in cp.process_data()}
+	return id2doc
+
+
+
+def proc_topics(topic_path: str, output_path: str, trec_or_kz: str = 'trec') -> Dict[str, CTTopic]:
+	
+	ct_config = CTConfig(
+		data_path=topic_path, 
+		write_file=output_path,
+    	nlp=True,
+    	is_topic=True,
+		trec_or_kz=trec_or_kz
+	)
+
+	cp = CTProc(ct_config)
+	id2topic = {res.id : res for res in cp.process_data()}
+	return id2topic
+
+
+
+def filter_doc_for_ir(doc, dconfig) -> Dict[str, List[str]]:
+	new_doc = dict()
+	new_doc['id'] = doc['id']
+	new_doc['text'] = prep_doc_text(doc, dconfig)
+	return new_doc
+
+
+def prep_ir_dataset(dconfig: DataConfig):
+	# need a file of all docs with their 
+	# 1. ids, 
+	# 2. combined text... 
+	# 3. 
+	
+	# get path to processed docs
+	doc_tuples, _ = ctpaths.get_data_tuples(dconfig.trec_or_kz)
+		
+	# get all processed docs
+	id2doc = dict()
+	for _, processed_doc_path in doc_tuples:
+		print(f"getting docs from {processed_doc_path}")
+		for doc in get_processed_data(processed_doc_path):
+			doc = filter_doc_for_ir(doc, dconfig)
+			doc['category'] = np.asarray(sorted(doc['category']).values())  # makes a consistently ordered category vector
+			id2doc[doc.id] = doc
+	return id2doc
+
+
+# --------------------------------------------------------------------------------------------------------------- #
+# pre-processing functions to save a form of triples for a particular model spec
+# --------------------------------------------------------------------------------------------------------------- #
+
+def prep_fine_tuning_dataset(
+		dconfig: DataConfig
+) -> None:
+	"""
+	trec_or_kz: 'trec' or 'kz'
+	desc: create dict of triplets of topic, doc, relevancy scores,
+	      save into a single jsonl file
+	"""
+	print(f"trec_or_kz: {dconfig.trec_or_kz}")
+	topic_path, rel_path = get_topic_and_rel_path(dconfig.trec_or_kz)
+	
+
+	# get set of all relevant doc ids
+	rel_type_dict, rel_dict, all_qrelled_docs = analyze_test_rels(rel_path)
+	
+	# get path to processed docs (already got topic path)
+	doc_tuples, _ = ctpaths.get_data_tuples(dconfig.trec_or_kz)
+
+	# get mappings of doc ids to doc dicts and topic ids to topic dicts
+	id2doc, id2topic = get_doc_and_topic_mappings(all_qrelled_docs, doc_tuples, topic_path) 
+	print(len(id2doc), len(all_qrelled_docs))
+	
+	missing_docs = set()
+	skipped = 0
+
+	# save combined triples of doc, topic, relevancy score
+	with open(dconfig.save_path, 'w') as f:
+		print(f"saving to: {dconfig.save_path}")
+
+		for topic_id in rel_dict:
+			for doc_id in rel_dict[topic_id]:
+				label = rel_dict[topic_id][doc_id]	
+				if downsample_zero(label, rel_type_dict['0'], dconfig):
+					skipped += 1
+					continue
+					
+				if doc_id in id2doc:
+					combined = create_combined_doc(
+						id2doc[doc_id], 
+						id2topic[topic_id], 
+						label, 
+						dconfig=dconfig, 
+					)
+
+					# save to file as jsonl
+					f.write(json.dumps(combined))
+					f.write('\n')
+				else:
+					missing_docs.add(doc_id)
+
+
+	print(f"number of docs missing: {len(missing_docs)}, number of zeros skipped: {skipped}")
+	for md in missing_docs:
+		print(md)
+
+
+
+def create_combined_doc(
+	doc, topic, 
+	rel_score, 
+	dconfig: DataConfig,
+):
+	combined = dict()
+
+	# get filtered and truncated and SEP tokenized topic text
+	combined['topic'] = prep_topic_text(topic, dconfig)
+
+	# get filtered and truncated and SEP tokenized doc text
+	combined['doc'] = prep_doc_text(doc, dconfig)
+
+	# get relevancy score as string 
+	if dconfig.convert_snli:
+		rel_score = convert_label_snli(rel_score)
+
+	combined['label'] = str(rel_score)
+
+	return combined
+
+
+def convert_label_snli(label: int) -> int:
+	if label == 2:
+		return 1
+	elif label == 1:
+		return 2
+	return label
+
+
+
+def downsample_zero(label: str, zero_ct: int, dconfig: DataConfig) -> bool:
+	if dconfig.downsample_zeros_n is not None:
+		if (label == 0) and (random.random()  >  (dconfig.downsample_zeros_n / zero_ct)):
+			return True
+	return False
+
+
+def prep_topic_text(topic: Dict[str, Union[List[str], str, float]], dconfig: DataConfig) -> str:
+	topic_text = ' '.join(topic['text_sents'])
+	topic_text = truncate(topic_text, dconfig.max_topic_len)
+	return topic_text
+
+
+def get_n_crit(crit_list: List[str], dconfig: DataConfig) -> List[str]:
+	if dconfig.first_n_only is not None:
+		crit_list = crit_list[:min(len(crit_list), dconfig.first_n_only)]
+	return crit_list
+
+
+def prep_doc_text(doc: Dict[str, Union[List[str], str, float]], dconfig: DataConfig) -> str:
+	
+	# combine lists of strings into single string
+	doc_inc = ' '.join(get_n_crit(doc['elig_crit']['include_criteria'], dconfig))
+	doc_exc = ' '.join(get_n_crit(doc['elig_crit']['exclude_criteria'], dconfig))
+
+
+	if 'condition' in dconfig.filtered_doc_keys:
+		doc_inc = f"{' '.join(doc['condition'])} {doc_inc}"
+		if dconfig.llm_prep:
+			doc_inc = "Condition: " + doc_inc + ", "
+
+	#truncate criteria separately if in config
+	doc_inc = truncate(doc_inc, dconfig.max_inc_len)
+	doc_exc = truncate(doc_exc, dconfig.max_exc_len)
+
+
+	if dconfig.prepend_elig_gender:
+		doc_inc = f"{doc['elig_gender']} {dconfig.sep} {doc_inc}"
+		if dconfig.llm_prep:
+			doc_inc = "Gender: " + doc_inc + ", "
+
+	if dconfig.prepend_elig_age:
+		if dconfig.llm_prep:
+			doc_inc = f"Trial Doc: A person who is between {doc['elig_min_age']}-{doc['elig_max_age']} years old who meets the following Inclusion Criteria: {doc_inc}"
+		else:
+			doc_inc = f"eligible ages (years): {doc['elig_min_age']}-{doc['elig_max_age']}, {dconfig.sep} {doc_inc}"
+
+	# combine criteria into single string
+	if dconfig.include_only:
+		if dconfig.llm_prep:
+			doc_inc += LLM_END_PROMPT
+		return doc_inc
+	
+	if dconfig.llm_prep:
+		return f"{doc_inc} and does not meet these Exclusion Criteria: {doc_exc} {LLM_END_PROMPT}"
+
+	return f"{doc_inc} {dconfig.sep} {doc_exc}"
+
+	
+
+
+# --------------------------------------------------------------------------------------------------------------- #
+# utility functions
+# --------------------------------------------------------------------------------------------------------------- #
+
+def age_match(min_doc_age: float, max_doc_age: float, topic_age: float) -> bool:
+	if topic_age < min_doc_age:
+		return False
+	if topic_age > max_doc_age:
+		return False
+	return True
+
+def gender_match(doc_gender: str, topic_gender: str) -> bool:
+	if doc_gender == 'All':
+		return True
+	if doc_gender == topic_gender:
+		return True
+	return False
+
+
+def get_topic_and_rel_path(trec_or_kz: str = 'trec') -> Tuple[str, str]:
+	if trec_or_kz == 'trec':
+		rel_path = ctpaths.TREC_REL_PATH
+		topic_path = ctpaths.TREC_RELLED_TOPIC_PATH
+	else:
+		rel_path = ctpaths.KZ_REL_PATH
+		topic_path = ctpaths.KZ_RELLED_TOPIC_PATH
+	return topic_path, rel_path
+
+
+def get_doc_and_topic_mappings(all_qrelled_docs: Set[str], doc_tuples: List[Tuple[str, str]], topic_path: str) -> Tuple[Dict[str, Dict[str, str]], Dict[str, Dict[str, str]]]:
+	"""
+	desc: get mappings of doc ids to doc dicts and topic ids to topic dicts
+	"""
+	
+	# get all processed topics
+	id2topic = {t['id']:t for t in get_processed_data(topic_path)}
+
+	# get all processed docs
+	id2doc = dict()
+	for _, processed_doc_path in doc_tuples:
+		print(f"getting docs from {processed_doc_path}")
+		for doc in get_processed_data(processed_doc_path):
+			if doc['id'] in all_qrelled_docs:
+				id2doc[doc['id']] = doc
+	
+	return id2doc, id2topic
+
+
+if __name__ == '__main__':
+	# proc_docs_and_topics('kz')
+	# eda.explore_trec_data(part=2, rand_print=0.001) # select part 1-5 (~70k docs per part)
+	# eda.explore_kz_data(rand_print=0.00001) # all in one file (~200k docs)
+
+# example config:
+# class DataConfig(NamedTuple):
+# 	save_path: str
+# 	trec_or_kz: str = 'trec'
+# 	filtered_topic_keys: Set[str] = {'id', 'text_sents', 'age', 'gender'}
+# 	filtered_doc_keys: Set[str] = {'id', 'elig_min_age', 'elig_max_age', 'elig_gender', 'condition', 'elig_crit'}
+# 	max_topic_len: Optional[int] = None
+# 	max_inc_len: Optional[int] = None
+# 	max_exc_len: Optional[int] = None
+# 	prepend_elig_age: bool = True
+# 	prepend_elig_gender: bool = True
+# 	include_only: bool = False
+# 	downsample_zeros_n: Optional[int] = None
+# 	sep: str = '[SEP]'
+# 	llm_prep: bool = False
+# 	first_n_only: Optional[int] = None
+# 	convert_snli: bool = False
+# 	infer_category_model: Optional[str] = None
+
+	dconfig = DataConfig(
+		trec_or_kz='trec',
+		save_path=ctpaths.TREC_ML_PATH, # make sure to change this!
+		sep='',
+		first_n_only=10,
+		max_topic_len=200,
+		llm_prep=False,
+		prepend_elig_age=True,
+		prepend_elig_gender=False
+	)
+	prep_fine_tuning_dataset(dconfig)
+	#eda.explore_prepped(ctpaths.TREC_KZ_PATH)
+

ctmatch/dataprep.py

@@ -0,0 +1,152 @@
+
+
+
+# external imports
+from datasets import Dataset, load_dataset, ClassLabel, Features, Value
+from transformers import AutoTokenizer
+import pandas as pd
+import numpy as np
+
+# package tools
+from .utils.ctmatch_utils import train_test_val_split, get_processed_data, get_test_rels
+from .pipeconfig import PipeConfig
+
+
+# path to ctmatch dataset on HF hub
+CTMATCH_CLASSIFICATION_DATASET_ROOT = "semaj83/ctmatch_classification"
+CTMATCH_IR_DATASET_ROOT = "semaj83/ctmatch_ir"
+CLASSIFIER_DATA_PATH = "combined_classifier_data.jsonl"
+DOC_TEXTS_PATH = "doc_texts.txt"
+DOC_CATEGORIES_VEC_PATH = "doc_categories.txt"
+DOC_EMBEDDINGS_VEC_PATH = "doc_embeddings.txt"
+INDEX2DOCID_PATH = "index2docid.txt"
+
+
+SUPPORTED_LMS = [
+    'roberta-large', 'cross-encoder/nli-roberta-base',
+    'microsoft/biogpt', 'allenai/scibert_scivocab_uncased', 
+    'facebook/bart-large', 'gpt2', 
+    'semaj83/scibert_finetuned_ctmatch', 'semaj83/scibert_finetuned_pruned_ctmatch'
+
+]
+
+
+class DataPrep:
+    # multiple 'datasets' need to be prepared for the pipeline
+    # 1. the dataset for the classifier model triplets and a dataframe, ~ 25k rows
+    # 2. the dataset for the category model, every doc ~200k rows
+    # 3. the dataset for the embedding model, every doc < 200k rows
+
+
+
+    def __init__(self, pipe_config: PipeConfig) -> None:
+        self.pipe_config = pipe_config
+        self.classifier_tokenizer = self.get_classifier_tokenizer()
+        self.ct_dataset = None
+        self.ct_train_dataset_df = None
+        self.index2docid = None
+        self.doc_embeddings_df = None
+        self.doc_categories_df = None
+
+        if pipe_config.ir_setup:
+            self.load_ir_data()
+        else:
+            self.load_classifier_data()
+
+
+        
+	
+    def get_classifier_tokenizer(self):
+        model_checkpoint = self.pipe_config.classifier_model_checkpoint
+        if model_checkpoint not in SUPPORTED_LMS:
+            raise ValueError(f"Model checkpoint {model_checkpoint} not supported. Please use one of {SUPPORTED_LMS}")
+        if 'scibert' in model_checkpoint:
+            tokenizer = AutoTokenizer.from_pretrained('allenai/scibert_scivocab_uncased', use_fast=True)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(self.pipe_config.classifier_model_checkpoint)
+        if self.pipe_config.classifier_model_checkpoint == 'gpt2':
+            tokenizer.pad_token = tokenizer.eos_token
+        return tokenizer
+
+
+    # ------------------ Classifier Data Loading ------------------ #
+    def load_classifier_data(self) -> Dataset:
+        self.ct_dataset = load_dataset(CTMATCH_CLASSIFICATION_DATASET_ROOT, data_files=CLASSIFIER_DATA_PATH)
+        self.ct_dataset = train_test_val_split(self.ct_dataset, self.pipe_config.splits, self.pipe_config.seed)
+        self.add_features()
+        self.tokenize_dataset()
+        self.ct_dataset = self.ct_dataset.rename_column("label", "labels")
+        # self.ct_dataset = self.ct_dataset.rename_column("topic", "sentence1")
+        # self.ct_dataset = self.ct_dataset.rename_column("doc", "sentence2")
+        self.ct_dataset.set_format(type='torch', columns=['doc', 'labels', 'topic', 'input_ids', 'attention_mask'])
+        if not self.pipe_config.use_trainer:
+            self.ct_dataset = self.ct_dataset.remove_columns(['doc', 'topic'])  # removing labels for next-token prediction...
+
+        self.ct_train_dataset_df = self.ct_dataset['train'].remove_columns(['input_ids', 'attention_mask', 'token_type_ids']).to_pandas()
+
+        return self.ct_dataset
+
+    
+    def add_features(self) -> None:
+        if self.pipe_config.convert_snli:
+            names = ['contradiction', 'entailment', 'neutral']
+        else:
+            names = ["not_relevant", "partially_relevant", "relevant"]
+
+        features = Features({
+            'doc': Value(dtype='string', id=None),
+            'label': ClassLabel(names=names),
+            'topic': Value(dtype='string', id=None)
+        })
+        self.ct_dataset["train"] = self.ct_dataset["train"].map(lambda x: x, batched=True, features=features)
+        self.ct_dataset["test"] = self.ct_dataset["test"].map(lambda x: x, batched=True, features=features)
+        self.ct_dataset["validation"] = self.ct_dataset["validation"].map(lambda x: x, batched=True, features=features)  
+
+
+    def tokenize_function(self, examples):
+        return self.classifier_tokenizer(
+            examples["topic"], examples["doc"], 
+            truncation=self.pipe_config.truncation, 
+            padding=self.pipe_config.padding, 
+            max_length=self.pipe_config.max_length
+        )
+
+    def tokenize_dataset(self):
+        self.ct_dataset = self.ct_dataset.map(self.tokenize_function, batched=True)
+
+
+    def get_category_data(self, vectorize=True):
+        category_data = dict()
+        sorted_cat_keys = None
+        for cdata in get_processed_data(self.pipe_config.category_path):
+
+            # cdata = {<nct_id>: {cat1: float1, cat2: float2...}}
+            cdata_id, cdata_dict = list(cdata.items())[0]
+            if sorted_cat_keys is None:
+                sorted_cat_keys = sorted(cdata_dict.keys())
+
+            if vectorize:
+                cat_vec = np.asarray([cdata_dict[k] for k in sorted_cat_keys])
+            else:
+                cat_vec = cdata_dict
+
+            category_data[cdata_id] = cat_vec
+        return category_data
+    
+
+    
+    # ------------------ IR Data Loading ------------------ #
+    def process_ir_data_from_hf(self, ds_path, is_text: bool = False):
+        ds = load_dataset(CTMATCH_IR_DATASET_ROOT, data_files=ds_path)
+        if is_text:
+            return pd.DataFrame(ds['train'])
+    
+        arrays = [np.asarray(a['text'].split(','), dtype=float) for a in ds['train']]
+        return pd.DataFrame(arrays)
+
+    def load_ir_data(self) -> None:
+        self.index2docid = self.process_ir_data_from_hf(INDEX2DOCID_PATH, is_text=True)
+        self.doc_embeddings_df = self.process_ir_data_from_hf(DOC_EMBEDDINGS_VEC_PATH)
+        self.doc_categories_df = self.process_ir_data_from_hf(DOC_CATEGORIES_VEC_PATH)
+        self.doc_texts_df = self.process_ir_data_from_hf(DOC_TEXTS_PATH, is_text=True)
+

ctmatch/eda.py

@@ -0,0 +1,114 @@
+	
+from typing import Dict, NamedTuple, Tuple
+from utils.ctmatch_utils import get_processed_data
+from collections import defaultdict
+import ct_data_paths
+import random
+
+from ctproc.scripts.vis_scripts import (
+  analyze_test_rels
+)
+
+
+class ExplorePaths(NamedTuple):
+	doc_path: str
+	topic_path: str
+	rel_path: str
+
+
+
+# --------------------------------------------------------------------------------------------------------------- #
+# EDA functions
+# --------------------------------------------------------------------------------------------------------------- #
+
+def explore_kz_data(rand_print: float = 0.001) -> None:
+	kz_data_paths = ExplorePaths(
+		rel_path = ct_data_paths.KZ_REL_PATH,
+		doc_path = ct_data_paths.KZ_PROCESSED_DOC_PATH,
+		topic_path = ct_data_paths.KZ_RELLED_TOPIC_PATH
+	)
+	
+	explore_data(kz_data_paths, rand_print=rand_print)
+
+
+def explore_trec_data(part: int = 1, rand_print: float = 0.001) -> None:
+	# post processing analysis
+	trec_data_paths = ExplorePaths(
+		rel_path = ct_data_paths.TREC_REL_PATH,
+		doc_path = f'/Users/jameskelly/Documents/cp/ctmatch/data/trec_data/processed_trec_data/processed_trec22_docs_part{part}.jsonl',
+		topic_path = ct_data_paths.TREC_RELLED_TOPIC_PATH
+	)
+	
+	explore_data(trec_data_paths, rand_print=rand_print)
+
+
+
+def explore_data(data_paths: ct_data_paths.ExplorePaths, rand_print: float) -> None:
+
+	# process relevancy judgements
+	type_dict, rel_dict, all_qrelled_docs = analyze_test_rels(data_paths.rel_path)
+	
+	# get processed topics
+	id2topic = {t['id']:t for t in get_processed_data(data_paths.topic_path)}
+	print(f"number of processed topics: {len(id2topic)}")
+
+	# get relevant processed docs
+	id2docs = {doc['id']:doc for doc in get_processed_data(data_paths.doc_path, get_only=all_qrelled_docs)}
+	print(f"number of relevant processed docs: {len(id2docs)}")
+
+	explore_pairs(id2topic, id2docs, rel_dict, max_print=1000, rand_print=rand_print)
+
+
+
+
+
+def explore_pairs(id2topic: Dict[str, Dict[str, str]], id2docs: Dict[str, Dict[str, str]], rel_dict: Dict[str, Dict[str, str]], rand_print: float, max_print:int = 100000) -> None:
+	rel_scores = defaultdict(int)
+	age_mismatches, gender_mismatches = 0, 0
+	for pt_id, topic in id2topic.items():
+		for doc_id in rel_dict[pt_id]:
+			if doc_id in id2docs:
+				rel_score = rel_dict[pt_id][doc_id]
+				rel_scores[rel_score] += 1
+				if rel_score == 2:
+					age_mismatches, gender_mismatches = check_match(
+						topic = topic, 
+						doc = id2docs[doc_id], 
+						rel_score = rel_score,
+						age_mismatches = age_mismatches, 
+						gender_mismatches = gender_mismatches
+					)
+
+				if random.random() < rand_print:
+					print_pair(topic, id2docs[doc_id], rel_score, marker='%')
+
+	print(rel_scores.items())
+	print(f"{age_mismatches=}, {gender_mismatches=}")
+
+
+
+
+def check_match(topic: Dict[str, str], doc: Dict[str, str], rel_score: int, age_mismatches: int, gender_mismatches: int) -> Tuple[int, int]:
+	age_matches = age_match(doc['elig_min_age'], doc['elig_max_age'], topic['age'])
+	if not age_matches:
+		#print_pair(topic, doc, rel_score)
+		age_mismatches += 1
+
+	gender_matches = gender_match(doc['elig_gender'], topic['gender'])
+	if not gender_matches:
+		#print_pair(topic, doc, rel_score)
+		gender_mismatches += 1
+	
+	return age_mismatches, gender_mismatches
+
+
+					
+def print_pair(topic: Dict[str, str], doc: Dict[str, str], rel_score: int, marker: str = '*') -> None:
+	print(marker*200)
+	print(f"topic id: {topic['id']}, nct_id: {doc['id']}, rel score: {rel_score}")
+	print(f"topic info: \nage: {topic['age']}, gender: {topic['gender']}")
+	print(topic['raw_text'])
+	print(f"doc info: gender: {doc['elig_gender']}, min age: {doc['elig_min_age']}, max age: {doc['elig_max_age']}")
+	print(doc['elig_crit']['raw_text'])
+	print(marker*200)
+	print()

ctmatch/evaluator.py

@@ -0,0 +1,154 @@
+
+import logging
+from typing import List, NamedTuple, Optional, Tuple, Union
+
+from .utils.eval_utils import (
+    calc_first_positive_rank, calc_f1, get_kz_topic2text, get_trec_topic2text
+)
+from .pipeconfig import PipeConfig
+from .match import CTMatch
+from pathlib import Path
+from tqdm import tqdm
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class EvaluatorConfig(NamedTuple):
+    rel_paths: List[str]
+    trec_topic_path: Union[Path, str]  = None
+    kz_topic_path: Union[Path, str] = None
+    max_topics: int = 200
+    openai_api_key: Optional[str] = None
+    filters: Optional[List[str]] = None
+    sanity_check_ids: Optional[List[str]] = None
+
+
+class Evaluator:
+    def __init__(self, eval_config: EvaluatorConfig) -> None:
+        self.rel_paths: List[str] = eval_config.rel_paths
+        self.trec_topic_path: Union[Path, str]  = eval_config.trec_topic_path
+        self.kz_topic_path: Union[Path, str] = eval_config.kz_topic_path
+       
+        self.rel_dict: dict = None
+        self.topicid2text: dict = None
+        self.ctm = None
+        self.openai_api_key = eval_config.openai_api_key
+        self.filters = eval_config.filters
+        self.sanity_check_ids = eval_config.sanity_check_ids
+   
+        assert self.rel_paths is not None, "paths to relevancy judgments must be set in pipe_config if pipe_config.evaluate=True"
+        assert ((self.trec_topic_path is not None) or (self.kz_topic_path is not None)), "at least one of trec_topic_path or kz_topic_path) must be set as pipe_config.evaluate=True"
+    
+        self.setup()
+
+        self.max_topics: int = len(self.topicid2text) if eval_config.max_topics is None else min(len(self.topicid2text), eval_config.max_topics)
+
+
+
+    def get_combined_rel_dict(self, rel_paths: List[str]) -> dict:
+        combined_rel_dict = dict()
+        for rel_path in rel_paths:
+            with open(rel_path, 'r') as f:
+                for line in f.readlines():
+                    topic_id, _, doc_id, rel = line.split()
+                    if topic_id not in combined_rel_dict:
+                        combined_rel_dict[topic_id] = dict()
+                    combined_rel_dict[topic_id][doc_id] = int(rel)
+        return combined_rel_dict
+    
+    def setup(self):
+        self.rel_dict = self.get_combined_rel_dict(self.rel_paths)
+        self.topicid2text = dict()
+        if self.kz_topic_path is not None:
+            self.topicid2text = get_kz_topic2text(self.kz_topic_path)
+
+        if self.trec_topic_path is not None:
+            self.topicid2text.update(get_trec_topic2text(self.trec_topic_path))
+
+        # loads all remaining needed datasets into memory
+        pipe_config = PipeConfig(
+            openai_api_key=self.openai_api_key,
+            ir_setup=True,
+            filters=self.filters
+        )
+        self.ctm = CTMatch(pipe_config=pipe_config)
+
+
+
+    def evaluate(self):
+        """
+        desc: run the pipeline over every topic and associated labelled set of documents,
+                and compute the mean mrr over all topics (how far down to the first relevant document)
+        """
+        frrs, f1s, fprs = [], [], []
+        for topic_id, topic_text in tqdm(list(self.topicid2text.items())[:self.max_topics]):
+       
+            if topic_id not in self.rel_dict:
+                # can't evaluate with no judgments
+                continue
+                
+            doc_ids = list(self.rel_dict[topic_id].keys())
+            logger.info(f"number of ranked docs: {len(doc_ids)}")
+            doc_set = self.get_indexes_from_ids(doc_ids)
+
+            # run IR pipeline on set of indexes corresponding to labelled doc_ids
+            ranked_pairs = self.ctm.match_pipeline(topic_text, doc_set=doc_set)
+
+            # get NCTIDs from ranking
+            ranked_ids = [nct_id for nct_id, doc_text in ranked_pairs]
+
+            # calculate metrics
+            fpr, frr = calc_first_positive_rank(ranked_ids, self.rel_dict[topic_id])
+            f1 = calc_f1(ranked_ids, self.rel_dict[topic_id])
+
+            if self.sanity_check_ids is not None and (topic_id in self.sanity_check_ids):
+                self.sanity_check(topic_id, topic_text, ranked_pairs, self.rel_dict[topic_id])
+
+            fprs.append(fpr)
+            frrs.append(frr)
+            f1s.append(f1)
+        
+        mean_fpr = sum(fprs)/len(fprs)
+        std_fpr = np.std(fprs)
+        mean_frr = sum(frrs)/len(frrs)
+        std_frr = np.std(frrs)
+        mean_f1 = sum(f1s)/len(f1s)
+        std_f1 = np.std(f1s)
+
+        return {
+            "mean_fpr":mean_fpr, "std_fpr":std_fpr,
+            "mean_frr":mean_frr, "std_frr":std_frr,
+            "mean_f1":mean_f1, "std_f1":std_f1
+        }
+
+
+    def get_indexes_from_ids(self, doc_id_set: List[str]) -> List[int]:
+        """
+        desc:       get the indexes of the documents in doc_id_set in the order they appear in the ranking
+        returns:    list of indexes
+        """
+        doc_indices = []
+        for doc_id in doc_id_set:
+            index_row = np.where(self.ctm.data.index2docid['text'] == doc_id)
+            if len(index_row[0]) == 0:
+                continue
+            doc_indices.append(index_row[0][0])
+        return doc_indices
+
+    def sanity_check(self, topic_id, topic_text, ranked_pairs: List[Tuple[str, str]], rel_dict) -> None:
+        logger.info(f"{topic_id=} {topic_text}")
+        for doc_id, doc_text in ranked_pairs:
+            rel_score = rel_dict[doc_id]
+            logger.info(rel_score, doc_id, doc_text)
+
+
+
+
+
+
+
+
+
+            
+        

ctmatch/match.py

@@ -0,0 +1,333 @@
+
+import logging
+from typing import Any, Dict, List, Optional, Tuple
+
+
+# external imports
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
+from numpy.linalg import norm
+from pathlib import Path
+from sklearn import svm
+import numpy as np
+import torch
+import json
+
+
+# package tools
+from .models.classifier_model import ClassifierModel
+from .utils.ctmatch_utils import get_processed_data, exclusive_argmax
+from .models.gen_model import GenModel
+from .pipeconfig import PipeConfig
+from .pipetopic import PipeTopic
+from .dataprep import DataPrep
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+CT_CATEGORIES = [
+    "pulmonary", "cardiac", "gastrointestinal", "renal", "psychological", "genetic", "pediatric",
+	"neurological", "cancer", "reproductive", "endocrine", "infection", "healthy", "other"
+]
+
+
+GEN_INIT_PROMPT =  "I will give you a patient description and a set of clinical trial documents. Each document will have a NCTID. I would like you to return the set of NCTIDs ranked from most to least relevant for patient in the description.\n"
+
+
+class CTMatch:
+    
+    def __init__(self, pipe_config: Optional[PipeConfig] = None) -> None:
+        # default to model config with full ir setup
+        self.pipe_config = pipe_config if pipe_config is not None else PipeConfig(ir_setup=True)
+        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        self.data = DataPrep(self.pipe_config)
+        self.classifier_model = ClassifierModel(self.pipe_config, self.data, self.device)
+        self.embedding_model = SentenceTransformer(self.pipe_config.embedding_model_checkpoint) 
+        self.gen_model = GenModel(self.pipe_config)
+        self.category_model = None
+        self.filters: Optional[List[str]] = pipe_config.filters
+
+        # filter params
+        self.sim_top_n = 10000
+        self.svm_top_n = 100
+        self.classifier_top_n = 50
+        self.gen_top_n = 10
+
+
+    # main api method
+    def match_pipeline(self, topic: str, top_k: int = 10, doc_set: Optional[List[int]] = None) -> List[str]:
+
+        if doc_set is None:
+            # start off will all doc indexes
+            doc_set = [i for i in range(len(self.data.index2docid))]
+        else:
+            self.reset_filter_params(len(doc_set))
+ 
+        # get topic representations for pipeline filters
+        pipe_topic = self.get_pipe_topic(topic)
+
+        if self.filters is None or ('sim' in self.filters):
+            # first filter, category + embedding similarity
+            doc_set = self.sim_filter(pipe_topic, doc_set, top_n=self.sim_top_n)
+
+        if self.filters is None or ('svm' in self.filters):
+            # second filter, SVM
+            doc_set = self.svm_filter(pipe_topic, doc_set, top_n=self.svm_top_n)
+      
+        if self.filters is None or ('classifier' in self.filters):
+            # third filter, classifier-LM (reranking)
+            doc_set = self.classifier_filter(pipe_topic, doc_set, top_n=self.classifier_top_n)
+
+        if self.filters is None or ('gen' in self.filters):
+            # fourth filter, generative-LM
+            doc_set = self.gen_filter(pipe_topic, doc_set, top_n=top_k)
+
+        return self.get_return_data(doc_set[:min(top_k, len(doc_set))])
+
+
+    def reset_filter_params(self, val: int) -> None:
+        self.sim_top_n = self.svm_top_n = self.classifier_top_n = self.gen_top_n = val
+  
+
+    # ------------------------------------------------------------------------------------------ #
+    # filtering methods
+    # ------------------------------------------------------------------------------------------ #
+
+    def sim_filter(self, pipe_topic: PipeTopic, doc_set: List[int], top_n: int) -> List[int]:
+        """
+        filter documents by similarity to topic
+        doing this with loop and cosine similarity instead of linear kernel because of memory issues
+        """
+        logger.info(f"running sim filter on {len(doc_set)} docs")
+
+        topic_cat_vec = exclusive_argmax(pipe_topic.category_vec)
+        norm_topic_emb = norm(pipe_topic.embedding_vec)
+        cosine_dists = []
+        for doc_idx in doc_set:
+            doc_cat_vec = self.redist_other_category(self.data.doc_categories_df.iloc[doc_idx].values)
+        
+            # only consider strongest predicted category
+            doc_cat_vec = exclusive_argmax(doc_cat_vec)
+            doc_emb_vec = self.data.doc_embeddings_df.iloc[doc_idx].values
+
+            topic_argmax = np.argmax(topic_cat_vec)
+            doc_argmax = np.argmax(doc_cat_vec)
+            cat_dist = 0. if (topic_argmax == doc_argmax) else 1.
+            emb_dist = np.dot(pipe_topic.embedding_vec, doc_emb_vec) / (norm_topic_emb * norm(doc_emb_vec))
+            combined_dist = cat_dist + emb_dist
+            cosine_dists.append(combined_dist)
+
+        sorted_indices = list(np.argsort(cosine_dists))[:min(len(doc_set), top_n)]
+
+        # return top n doc indices by combined similiarity, biggest to smallest
+        return [doc_set[i] for i in sorted_indices]
+    
+
+    def svm_filter(self, topic: PipeTopic, doc_set: List[int], top_n: int) -> List[int]:
+        """
+           filter documents by training an SVM on topic and doc embeddings
+        """
+        logger.info(f"running svm filter on {len(doc_set)} documents")
+
+        # build training data and prediction vector of single positive class for SVM
+        topic_embedding_vec = topic.embedding_vec[np.newaxis, :]
+        x = np.concatenate([topic_embedding_vec, self.data.doc_embeddings_df.iloc[doc_set].values], axis=0)
+        y = np.zeros(len(doc_set) + 1)
+        y[0] = 1
+
+        # define and fit SVM
+        clf = svm.LinearSVC(class_weight='balanced', verbose=False, max_iter=10000, tol=1e-6, C=0.1)
+        clf.fit(x, y) 
+        
+        # infer for similarities
+        similarities = clf.decision_function(x)
+
+        # get top n doc indices by similiarity, biggest to smallest
+        result = list(np.argsort(-similarities)[:min(len(doc_set) + 1, top_n + 1)])
+
+        # remove topic from result
+        result.remove(0)
+
+        # indexes got shifted by 1 because topic was included in doc_set
+        return [doc_set[(r - 1)] for r in result]
+
+
+
+    def classifier_filter(self, pipe_topic: PipeTopic, doc_set: List[int], top_n: int) -> List[int]:
+        """
+        filter documents by classifier no relevance prediction
+        """
+        logger.info(f"running classifier filter on {len(doc_set)} documents")
+
+        # get doc texts
+        doc_texts = [v[0] for v in self.data.doc_texts_df.iloc[doc_set].values]
+    
+        # sort by reverse irrelevant prediction
+        neg_predictions = np.asarray([p[0] for p in self.classifier_model.batch_inference(pipe_topic.topic_text, doc_texts, return_preds=True)])
+       
+        # return top n doc indices by classifier, biggest to smallest
+        sorted_indices = list(np.argsort(neg_predictions)[:min(len(doc_set), top_n)])
+        return [doc_set[i] for i in sorted_indices]
+       
+
+
+    def gen_filter(self, topic: PipeTopic, doc_set: List[int], top_n: int = 10) -> List[int]:
+        """
+            gen model supplies a ranking of remaming docs by evaluating the pairs of topic and doc texts
+
+            in order to overcome the context length limitation, we need to do a kind of left-binary search over multiple 
+            prompts to arrive at a ranking that meets the number of documents requirement (top_n)
+
+            may take a few minutes to run through all queries and subqueries depending on size of doc_set
+        
+        """
+        logger.info(f"running gen filter on {len(doc_set)} documents")
+
+        assert top_n > 0, "top_n must be greater than 0"
+
+        ranked_docs = doc_set
+        iters = 0
+        while (len(ranked_docs) > top_n) and (iters < 10) and (len(ranked_docs) // 2 > top_n):
+            query_prompts = self.get_subqueries(topic, ranked_docs)
+
+            logger.info(f"calling gen model on {len(query_prompts)} subqueries")
+
+            # get gen model response for each query_prompt
+            subrankings = []
+            for prompt in query_prompts:
+                subrank = self.gen_model.gen_response(prompt)
+
+                # keep the top half of each subranking
+                subrankings.extend(subrank[:len(subrank) // 2])
+
+            ranked_docs = subrankings
+            iters += 1
+        
+        return ranked_docs[:min(len(ranked_docs), top_n)]
+
+    # ------------------------------------------------------------------------------------------ #
+    # filter helper methods
+    # ------------------------------------------------------------------------------------------ #
+
+    def get_pipe_topic(self, topic):
+        pipe_topic = PipeTopic(
+            topic_text=topic, 
+            embedding_vec=self.get_embeddings([topic])[0],             # 1 x embedding_dim (default=384)
+            category_vec=self.get_categories(topic)                    # 1 x 14
+        )
+        return pipe_topic
+    
+
+    def get_embeddings(self, texts: List[str]) -> List[float]:
+        return self.embedding_model.encode(texts)
+    
+    def get_categories(self, text: str) -> str:
+        if self.category_model is None:
+            self.category_model = pipeline(
+                'zero-shot-classification', 
+                model=self.pipe_config.category_model_checkpoint, 
+                device=0
+            )
+        output = self.category_model(text, candidate_labels=CT_CATEGORIES)
+        score_dict = {output['labels'][i]:output['scores'][i] for i in range(len(output['labels']))}
+
+        # to be consistent with doc category vecs 
+        sorted_keys = sorted(score_dict.keys())
+        return self.redist_other_category(np.array([score_dict[k] for k in sorted_keys]))
+
+    def redist_other_category(self, category_vec: np.ndarray, other_dim:int = 8) -> np.ndarray:
+        """
+            redistribute 'other' category weight to all other categories
+        """
+        other_wt = category_vec[other_dim] 
+        other_wt_dist = other_wt / (len(category_vec) - 1)
+        redist_cat_vec = category_vec + other_wt_dist
+        redist_cat_vec[other_dim] = 0
+        return redist_cat_vec
+        
+
+    def get_gen_query_prompt(self, topic: PipeTopic, doc_set: List[int]) -> str:
+        query_prompt = f"{GEN_INIT_PROMPT}Patient description: {topic.topic_text}\n"
+        
+        for i, doc_text in enumerate(self.data.doc_texts_df.iloc[doc_set].values):
+            query_prompt += f"NCTID: {doc_set[i]}, "
+            query_prompt += f"Eligbility Criteria: {doc_text[0]}\n"
+
+            # not really token length bc not tokenized yet but close enough if we undershoot
+            prompt_len = len(query_prompt.split()) 
+            if prompt_len > self.pipe_config.max_query_length:
+                break
+    
+        return query_prompt, i       
+
+
+    def get_subqueries(self, topic: PipeTopic, doc_set: List[int]) -> List[str]:
+        query_prompts = []
+        i = 0
+        while i < len(doc_set) - 1:
+
+            # break the querying over remaining doc set into multiple prompts
+            query_prompt, used_i = self.get_gen_query_prompt(topic, doc_set[i:])
+            query_prompts.append(query_prompt)
+            i += used_i
+
+        return query_prompts
+
+
+    def get_return_data(self, doc_set: List[int]) -> List[Tuple[str, str]]:
+        return_data = []
+        for idx in doc_set:
+            nctid = self.data.index2docid.iloc[idx].values[0]
+            return_data.append((nctid, self.data.doc_texts_df.iloc[idx].values[0]))
+        return return_data
+
+
+
+    # ------------------------------------------------------------------------------------------ #
+    # data prep methods that rely on model in CTMatch object (not run during routine program) 
+    # ------------------------------------------------------------------------------------------ #
+
+    def prep_ir_text(self, doc: Dict[str, List[str]], max_len: int = 512) -> str:
+        inc_text = ' '.join(doc['elig_crit']['include_criteria'])
+        exc_text = ' '.join(doc['elig_crit']['exclude_criteria'])
+        all_text = f"Inclusion Criteria: {inc_text}, Exclusion Criteria: {exc_text}"
+        split_text = all_text.split()
+        return ' '.join(split_text[:min(max_len, len(split_text))])
+
+
+    def prep_and_save_ir_dataset(self):
+        category_data = self.data.get_category_data()
+        with open(self.pipe_config.ir_save_path, 'w') as wf:
+            for ir_data in self.prep_ir_data():
+                ir_data['categories'] = str(category_data[ir_data['id']])
+                wf.write(json.dumps(ir_data))
+                wf.write('\n')
+
+
+    def prep_ir_data(self):
+        for data_path in self.pipe_config.processed_data_paths:
+            for i, doc in enumerate(get_processed_data(data_path)):
+                if i % 10000 == 0:
+                    logger.info(f"Prepping doc {i}")
+
+                ir_data_entry = dict()
+                ir_data_entry['id'] = doc['id']
+                doc_text = self.prep_ir_text(doc)
+                ir_data_entry['doc_text'] = doc_text
+                yield ir_data_entry
+
+
+    def save_texts(self) -> Dict[int, str]:
+        idx2id = dict()
+        with open(Path(self.pipe_config.ir_save_path).parent / 'texts', 'w', encoding='utf-8') as wf:
+            for i, doc in enumerate(get_processed_data(self.pipe_config.ir_save_path)):
+                idx2id[i] = doc['id']
+                if i % 10000 == 0:
+                    logger.info(f"Prepping doc {i}")
+
+                wf.write(doc['doc_text'])
+                wf.write('\n')
+        return idx2id
+    

ctmatch/models/classifier_model.py

@@ -0,0 +1,396 @@
+
+import logging
+from pathlib import Path
+from tqdm.auto import tqdm
+from typing import List, Tuple
+
+from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, get_scheduler
+from optimum.onnxruntime import ORTModelForSequenceClassification
+from optimum.onnxruntime.configuration import OptimizationConfig
+from optimum.onnxruntime import ORTOptimizer
+import evaluate
+
+from sklearn.metrics import confusion_matrix, classification_report
+from sklearn.metrics import f1_score
+from torch.utils.data import DataLoader
+from torch.optim import AdamW
+from torch import nn
+import torch
+
+from nn_pruning.patch_coordinator import ModelPatchingCoordinator, SparseTrainingArguments
+from nn_pruning.inference_model_patcher import optimize_model
+from nn_pruning.sparse_trainer import SparseTrainer
+
+
+from ..pipeconfig import PipeConfig
+from ..dataprep import DataPrep
+
+
+logger = logging.getLogger(__name__)
+
+PRUNED_HUB_MODEL_NAME = 'semaj83/scibert_finetuned_pruned_ctmatch'
+
+
+class WeightedLossTrainer(Trainer):
+    def __init__(self, label_weights, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.label_weights = label_weights
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        outputs = model(**inputs)
+        logits = outputs.get("logits")
+        labels = inputs.get("labels")
+        loss_func = nn.CrossEntropyLoss(weight=self.label_weights)
+        loss = loss_func(logits, labels)
+        return (loss, outputs) if return_outputs else loss
+
+
+
+
+class PruningTrainer(SparseTrainer, WeightedLossTrainer):
+    def __init__(self, sparse_args, *args, **kwargs):
+        WeightedLossTrainer.__init__(self, *args, **kwargs)
+        SparseTrainer.__init__(self, sparse_args)
+
+
+class ClassifierModel:
+    
+    def __init__(self, model_config: PipeConfig, data: DataPrep, device: str):
+        self.model_config = model_config
+        self.dataset = data.ct_dataset
+        self.tokenizer = data.classifier_tokenizer
+        self.tokenize_func = data.tokenize_function
+        self.trainer = None
+        self.optimizer = None
+        self.lr_scheduler = None
+        self.device = device
+
+        if not self.model_config.ir_setup:
+            self.train_dataset_df = data.ct_dataset['train'].to_pandas()
+            self.num_training_steps = self.model_config.train_epochs * len(self.dataset['train'])
+
+        self.model = self.load_model()
+        self.pruned_model = None
+    
+        if not self.model_config.use_trainer and not self.model_config.ir_setup:
+            self.train_dataloader, self.val_dataloader = self.get_dataloaders()
+
+        
+        if self.model_config.prune:
+             self.prune_trainer = None
+             self.sparse_args = self.get_sparse_args()
+             self.mpc = self.get_model_patching_coordinator()
+               
+
+    # ------------------ Model Loading ------------------ #
+    def get_model(self):
+        if self.model_config.num_classes == 0:
+            return AutoModelForSequenceClassification.from_pretrained(self.model_config.classifier_model_checkpoint)
+
+        id2label, label2id = self.get_label_mapping()
+        model = AutoModelForSequenceClassification.from_pretrained(
+            self.model_config.classifier_model_checkpoint,
+            num_labels=self.model_config.num_classes,     # makes the last head be replaced with a linear layer with num_labels outputs (fine-tuning)
+            id2label=id2label, label2id=label2id,
+            ignore_mismatched_sizes=True                  # because of pruned model changes
+        )
+
+        if 'pruned' in self.model_config.classifier_model_checkpoint:
+            model = optimize_model(model, "dense")
+
+        return self.add_pad_token(model)
+
+
+    def add_pad_token(self, model):
+        if model.config.pad_token_id is None:
+            model.config.pad_token_id = model.config.eos_token_id
+        return model
+
+
+    def load_model(self):
+        self.model = self.get_model()
+
+        if self.model_config.ir_setup:
+            return self.model
+        
+        self.optimizer = AdamW(self.model.parameters(), lr=self.model_config.learning_rate, weight_decay=self.model_config.weight_decay)
+        self.num_training_steps = self.model_config.train_epochs * len(self.dataset['train'])
+        self.lr_scheduler = get_scheduler(
+            name="linear", 
+            optimizer=self.optimizer, 
+            num_warmup_steps=self.model_config.warmup_steps, 
+            num_training_steps=self.num_training_steps
+        )
+
+        if self.model_config.use_trainer and not self.model_config.prune:
+            self.trainer = self.get_trainer()
+        else:
+            self.model = self.model.to(self.device)
+
+        return self.model
+
+
+    def get_label_mapping(self):
+        #id2label = {idx:self.dataset['train'].features["labels"].int2str(idx) for idx in range(3)}
+        id2label =  {'0':'not_relevant', '1':'partially_relevant', '2':'relevant'}
+        label2id = {v:k for k, v in id2label.items()}
+        return id2label, label2id
+
+    def get_label_weights(self):
+        label_weights = (1 - (self.train_dataset_df["labels"].value_counts().sort_index() / len(self.train_dataset_df))).values
+        label_weights = torch.from_numpy(label_weights).float().to("cuda")
+
+
+    def get_trainer(self):
+        return WeightedLossTrainer(
+            model=self.model,
+            optimizers=(self.optimizer, self.lr_scheduler),
+            args=self.get_training_args_obj(),
+            compute_metrics=self.compute_metrics,
+            train_dataset=self.dataset["train"],
+            eval_dataset=self.dataset["validation"],
+            tokenizer=self.tokenizer,
+            label_weights=self.get_label_weights()
+        )
+
+
+    def get_training_args_obj(self):
+        output_dir = self.model_config.output_dir if self.model_config.output_dir is not None else self.model_config.classifier_data_path.parent.parent.as_posix()
+        
+        return TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=self.model_config.train_epochs,
+            learning_rate=self.model_config.learning_rate,
+            per_device_train_batch_size=self.model_config.batch_size,
+            per_device_eval_batch_size=self.model_config.batch_size,
+            weight_decay=self.model_config.weight_decay,
+            evaluation_strategy="epoch",
+            logging_steps=len(self.dataset["train"]) // self.model_config.batch_size,
+            fp16=self.model_config.fp16
+        )
+        
+    
+
+    def train_and_predict(self):
+        if self.trainer is not None:
+            self.trainer.train()
+            predictions = self.trainer.predict(self.dataset["test"])
+            logger.info(predictions.metrics.items())
+        else:
+            self.loss_func = nn.CrossEntropyLoss(weight=self.get_label_weights())
+            self.manual_train()
+            self.manual_eval()
+
+
+
+     # ------------------ native torch training loop ------------------ #
+    def get_dataloaders(self) -> Tuple[DataLoader, DataLoader]:
+        train_dataloader = DataLoader(self.dataset['train'], shuffle=True, batch_size=self.model_config.batch_size)
+        val_dataloader = DataLoader(self.dataset['validation'], batch_size=self.model_config.batch_size)
+        return train_dataloader, val_dataloader
+
+
+
+    # taken from ctmatch for messing about 
+    def manual_train(self):
+        progress_bar = tqdm(range(self.num_training_steps))
+        self.model.train()
+        for epoch in range(self.model_config.train_epochs):
+            for batch in tqdm(self.train_dataloader):
+                batch = {k: v.to(self.model.device) for k, v in batch.items()}
+                outputs = self.model(**batch)
+                loss = self.loss_func(outputs.logits, batch['labels'])
+                #total_loss += loss.item()
+                loss.backward()
+                
+                self.optimizer.step()
+                self.lr_scheduler.step()
+                self.optimizer.zero_grad()
+
+                self.manual_eval()
+                logger.info(f"{loss=}")
+                progress_bar.update(1)
+            
+
+
+
+    def manual_eval(self):
+        metric = evaluate.load("f1")
+        self.model.eval()
+        for batch in self.val_dataloader:
+            batch = {k: v.to(self.model.device) for k, v in batch.items()}
+        
+            # don't learn during evaluation
+            with torch.no_grad():
+                outputs = self.model(**batch)
+
+            logits = outputs.logits
+            predictions = torch.argmax(logits, dim=-1)
+            metric.add_batch(predictions=predictions, references=batch["labels"])
+        
+        logger.info(metric.compute(average='weighted'))
+
+
+
+
+    def get_sklearn_metrics(self):
+        with torch.no_grad():
+            if self.model_config.use_trainer:
+                if self.model_config.prune:
+                    self.prune_trainer.model.to(self.device)
+                    logger.info("using pruned trainer model")
+                    preds = self.prune_trainer.predict(self.dataset['test']).predictions
+                else:    
+                    preds = self.trainer.predict(self.dataset['test']).predictions
+
+                if "bart" in self.model_config.name:
+                    preds = preds[0]
+
+                y_preds = list(preds.argmax(axis=1))
+            else:
+
+                if self.model_config.prune:
+                    model = self.pruned_model.to(self.device)
+                else:
+                    model = self.model.to(self.device)
+                y_preds = []
+                for input_ids in self.dataset['test']['input_ids']:
+                    input_ids = torch.tensor(input_ids).unsqueeze(0).to(self.device)
+                    y_pred = model(input_ids).logits.argmax().item()
+                    y_preds.append(y_pred)
+         
+        y_trues = list(self.dataset['test']['labels'])
+        return confusion_matrix(y_trues, y_preds), classification_report(y_trues, y_preds)
+    
+
+    def compute_metrics(self, pred):
+        labels = pred.label_ids
+        preds = pred.predictions
+        if "bart" in self.model_config.name:
+            preds = preds[0]
+        
+        preds = preds.argmax(-1)
+        f1 = f1_score(labels, preds, average="weighted")
+        return {"f1":f1}
+    
+    def inference_single_example(self, topic: str, doc: str, return_preds: bool = False) -> str:
+        """
+        desc: method to predict relevance label on new topic, doc examples 
+        """
+        ex = {'doc':doc, 'topic':topic}
+        with torch.no_grad():
+            inputs = torch.LongTensor(self.tokenize_func(ex)['input_ids']).unsqueeze(0)
+            outputs = self.model(inputs).logits
+            if return_preds:
+                return torch.nn.functional.softmax(outputs, dim=1).squeeze(0)
+            return str(outputs.argmax().item())
+
+
+    def batch_inference(self, topic: str, docs: List[str], return_preds: bool = False) -> List[str]:
+        topic_repeats = [topic for _ in range(len(docs))]
+        inputs = self.tokenizer(
+            topic_repeats, docs, return_tensors='pt', 
+            truncation=self.model_config.truncation, 
+            padding=self.model_config.padding, 
+            max_length=self.model_config.max_length
+        )
+
+        with torch.no_grad():
+            outputs = torch.nn.functional.softmax(self.model(**inputs).logits, dim=1)
+        
+        if return_preds:
+            return outputs
+        
+        return outputs.argmax(dim=1).tolist()
+
+
+
+    # ------------------ pruning  ------------------ #
+
+    def prune_model(self):
+        self.mpc.patch_model(self.model)
+        self.model.save_pretrained("models/patched")
+        self.prune_trainer = self.get_pruning_trainer()
+        self.prune_trainer.set_patch_coordinator(self.mpc)
+        self.prune_trainer.train()
+        self.mpc.compile_model(self.prune_trainer.model)
+        if self.model_config.push_to_hub:
+            # can't save the optimized model to hub
+            self.prune_trainer.model.push_to_hub(PRUNED_HUB_MODEL_NAME)
+
+        self.pruned_model = optimize_model(self.prune_trainer.model, "dense")
+      
+
+
+    def get_sparse_args(self):
+        sparse_args = SparseTrainingArguments()
+
+        hyperparams = {
+            "dense_pruning_method": "topK:1d_alt", 
+            "attention_pruning_method": "topK", 
+            "initial_threshold": 1.0, 
+            "final_threshold": 0.5, 
+            "initial_warmup": 1,
+            "final_warmup": 3,
+            "attention_block_rows":32,
+            "attention_block_cols":32,
+            "attention_output_with_dense": 0
+        }
+
+        for k,v in hyperparams.items():
+            if hasattr(sparse_args, k):
+                setattr(sparse_args, k, v)
+            else:
+                print(f"sparse_args does not have argument {k}")
+
+        return sparse_args
+
+    
+    def get_pruning_trainer(self):
+        return PruningTrainer(
+            sparse_args=self.sparse_args,
+            args=self.get_training_args_obj(),
+            model=self.model,
+            train_dataset=self.dataset["train"],
+            eval_dataset=self.dataset["validation"],
+            tokenizer=self.tokenizer,
+            compute_metrics=self.compute_metrics,
+            label_weights=self.get_label_weights()
+        )
+
+
+
+    def get_model_patching_coordinator(self):
+        return ModelPatchingCoordinator(
+            sparse_args=self.sparse_args, 
+            device=self.device, 
+            cache_dir="checkpoints", 
+            logit_names="logits", 
+            teacher_constructor=None
+        )
+    
+    
+    # onyx optimization
+    def optimize_model(self):  
+        onnx_path = Path("onnx")
+        model_id = self.model_config.classifier_model_checkpoint
+        #assert self.pruned_model is not None, "pruned model must be loaded before optimizing"
+        opt_model = ORTModelForSequenceClassification.from_pretrained(model_id, from_transformers=True)
+        optimizer = ORTOptimizer.from_pretrained(opt_model)
+        optimization_config = OptimizationConfig(optimization_level=99) # enable all optimizations
+        optimizer.optimize(
+            save_dir=onnx_path,
+            optimization_config=optimization_config,
+        )
+        opt_model.save_pretrained(onnx_path)
+        self.tokenizer.save_pretrained(onnx_path)  
+
+        #optimized_model = ORTModelForSequenceClassification.from_pretrained(onnx_path, file_name="model_optimized.onnx")  
+
+        return opt_model
+
+       
+
+
+
+        

ctmatch/models/gen_model.py

@@ -0,0 +1,83 @@
+
+from typing import List, Optional
+
+from ..pipeconfig import PipeConfig
+import openai
+import re
+
+
+
+
+class GenModel:
+    def __init__(self, pipe_config: PipeConfig) -> None:
+        openai.api_key = pipe_config.openai_api_key
+        self.pipe_config = pipe_config
+
+
+    def gen_response(self, query_prompt: str, doc_set: Optional[List[int]] = None) -> List[int]:
+        """
+        uses openai model to return a ranking of ids
+        """
+        if self.pipe_config.gen_model_checkpoint == 'text-davinci-003':
+            response = openai.Completion.create(
+                model=self.pipe_config.gen_model_checkpoint,
+                prompt=query_prompt,
+                temperature=0,
+                max_tokens=200,
+                top_p=1,
+                frequency_penalty=0.0,
+                presence_penalty=0.0
+            )
+        else:
+            assert doc_set is not None, "doc_set must be provided for gpt-3.5-turbo"
+            
+            # for gpt-3.5-turbo
+            response = openai.ChatCompletion.create(
+                model=self.pipe_config.gen_model_checkpoint,
+                messages = [{'role': 'user', 'content' : query_prompt}],
+                temperature=0.4,
+                max_tokens=200,
+                top_p=1,
+                frequency_penalty=0.2,
+                presence_penalty=0.0
+            )
+            
+
+        if self.pipe_config.gen_model_checkpoint == 'text-davinci-003':
+            return self.post_process_chatgpt_response(response)
+        return self.post_process_gptturbo_response(response, doc_set=doc_set)
+
+
+    def post_process_chatgpt_response(self, response):
+        """
+        could be:
+        NCTID 6, NCTID 7, NCTID 5
+        NCTID: 6, 7, 5
+        6, 7, 5
+        '1. 195155\n2. 186848\n3. 194407'
+        """
+        response_pattern = r"(?:NCTID\:?\s*)? ?(\d+)(?!\.)"
+        text = response['choices'][0]['text']
+        return [int(s) for s in re.findall(response_pattern, text)]
+
+    def post_process_gptturbo_response(self, response, doc_set: List[int]):
+        """
+        could be:
+        'The most relevant clinical trial for this patient is ID 2, followed by ID 3. The remaining trials are not relevant for this patient's condition.'
+        """
+        text = response['choices'][0]['message']['content']
+        ranking = []
+        for substr in text.split():
+            if substr.isdigit():
+                ranking.append(int(substr))
+
+        # the rest are arbitrarily ranked 
+        for ncid in doc_set:
+            if ncid not in ranking:
+                ranking.append(ncid)
+        return ranking
+
+
+
+
+	

ctmatch/pipeconfig.py

@@ -0,0 +1,42 @@
+
+from typing import Dict, List, NamedTuple, Optional
+from pathlib import Path
+
+
+class PipeConfig(NamedTuple):
+    name: str = 'scibert_finetuned_ctmatch'
+    classifier_model_checkpoint: str = 'semaj83/scibert_finetuned_ctmatch'
+    max_length: int = 512
+    padding: str = True
+    truncation: bool = True
+    batch_size: int = 16
+    learning_rate: float = 2e-5
+    train_epochs: int = 3
+    weight_decay: float = 0.01
+    warmup_steps: int = 500
+    seed: int  = 42
+    splits: Dict[str, float] = {"train":0.8, "val":0.1}
+    classifier_data_path: Path = Path("combined_classifier_data.jsonl")
+    output_dir: Optional[str] = None
+    convert_snli: bool = False
+    use_trainer: bool = False
+    num_classes: int = 3
+    fp16: bool = False
+    early_stopping: bool = False
+    push_to_hub: bool = False
+    ir_save_path: Optional[str] = None
+    category_path: Optional[str] = None       
+    processed_data_paths: Optional[List[str]] = None
+    max_query_length: int = 1200
+    category_model_checkpoint: str = "facebook/bart-large-mnli"
+    embedding_model_checkpoint: str = "sentence-transformers/all-MiniLM-L6-v2"
+    gen_model_checkpoint: str = 'text-davinci-003'
+    max_gen: int = 100
+    openai_api_key: Optional[str] = None
+    ir_setup: bool = False                  # if true, use the IR model setup, no classifier training or dataprep
+    filters: Optional[List[str]] = None     # if provided, only use these filters for the IR model, options are {'sim', 'svm', 'classifier', 'gen'}
+    prune: bool = False                     # if true, creates a pruned classifier model
+    optimize: bool = False                  # if true, creates an optimized classifier model
+
+
+

ctmatch/pipetopic.py

@@ -0,0 +1,10 @@
+
+from typing import Any, NamedTuple
+
+class PipeTopic(NamedTuple):
+    topic_text: str
+    embedding_vec: Any 
+    category_vec: Any
+  
+
+

ctmatch/scripts/build_combined_data.py

@@ -0,0 +1,76 @@
+
+from typing import Dict, List, Tuple
+import json
+
+
+COMBINED_CAT_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/combined_categories.jsonl'
+CAT_SAVE_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/doc_categories.txt'
+INDEX2DOCID_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/index2docid.txt'
+INDEX2TOPICID_PATH = '/Users/jameskelly/Documents/cp/ctmatch/data/index2topicid.txt'
+
+def load_category_dict(cat_path=COMBINED_CAT_PATH) -> Tuple[List, Dict[str, List[float]]]:
+    """
+    desc:   gets category dict from category path
+    """
+    sorted_cat_keys = None
+    
+    with open(cat_path, 'r') as json_file:
+        json_list = list(json_file)
+            
+    all_cat_dict = {}
+    for s in json_list:
+        s_data = json.loads(s)
+        nct_id, cat_dict = s_data.popitem()
+        
+        if sorted_cat_keys is None:
+            sorted_cat_keys = sorted(cat_dict.keys())
+
+        all_cat_dict[nct_id] = [cat_dict[k] for k in sorted_cat_keys]
+
+    return sorted_cat_keys, all_cat_dict
+
+
+
+def load_index2id(index2id_path: str = INDEX2DOCID_PATH) -> Dict[str, int]:
+    """
+    desc:   loads id2idx from csv path
+    """
+    index2id = {}
+    with open(index2id_path, 'r') as f:
+        for line in f:
+            if len(line) < 2:
+                continue
+            idx, nct_id = line.split(',')
+            index2id[idx] = nct_id.strip(' \n')
+    
+    return index2id
+
+
+
+def build_cat_csv(save_path: str = CAT_SAVE_PATH) -> None:
+    """
+    desc:   builds csv file for category data
+            VERY important that the indexes (order) match the order of the embeddings (for nctid lookup in idx2id)
+    """
+    sorted_cat_keys, cat_dict = load_category_dict()
+    idx2id = load_index2id()
+    
+    with open(save_path, 'w') as f:
+        f.write(','.join(sorted_cat_keys))
+        f.write('\n')
+        for _, nct_id in idx2id.items():
+            cat_vec = cat_dict[nct_id]
+            cat_vec_str = ','.join([str(c) for c in cat_vec])
+            f.write(cat_vec_str)
+            f.write('\n')
+
+
+if __name__ == '__main__':
+    build_cat_csv()
+
+    
+    
+    
+
+
+

ctmatch/scripts/gen_categories.py

@@ -0,0 +1,92 @@
+
+from typing import Generator, List, Optional, Tuple
+from ctmatch.utils.ctmatch_utils import get_processed_data
+from ctmatch.ct_data_paths import get_data_tuples
+from transformers import pipeline
+import numpy as np
+import json
+
+CAT_GEN_MODEL = "facebook/bart-large-mnli"
+#CAT_GEN_MODEL = "microsoft/biogpt"
+
+CT_CATEGORIES = [
+    "pulmonary", "cardiac", "gastrointestinal", "renal", "psychological", "genetic", "pediatric",
+	"neurological", "cancer", "reproductive", "endocrine", "infection", "healthy", "other"
+]
+
+
+# --------------------------------------------------------------------------------------------------------------- #
+# this script is for applying zero-shot classification labels from 'facebook/bart-large-mnli' to the 
+# documents of the dataset, including test, because we can assume this is something that is realistic to pre-compute
+# since you have the documents apriori
+# --------------------------------------------------------------------------------------------------------------- #
+GET_ONLY = None
+
+
+def stream_condition_data(data_chunk, doc_or_topic: str = 'doc') -> Generator[str, None, None]:
+    for d in data_chunk:
+      if doc_or_topic == 'topic':
+        yield d['raw_text']
+      else:
+        condition = d['condition']
+        if len(condition) == 0:
+          yield 'no information'
+        else:     
+          yield ' '.join(condition).lower()
+     
+
+def add_condition_category_labels(
+	trec_or_kz: str = 'trec', 
+	model_checkpoint=CAT_GEN_MODEL, 
+	start: int = 0,
+	doc_tuples: Optional[List[Tuple[str, str]]] = None,
+    category_label='category',
+    doc_or_topic: str = 'doc'
+) -> None:
+    pipe = pipeline(model=model_checkpoint, device=0)
+    chunk_size = 1000
+    
+    # open the processed documents and add the category labels
+    if doc_tuples is None:
+        doc_tuples, _ = get_data_tuples(trec_or_kz=trec_or_kz)
+        
+    for _, target in doc_tuples:
+        print(f"reading and writing to: {target}")
+        data = [d for d in get_processed_data(target, get_only=GET_ONLY)]
+        print(f"got {len(data)} records from {target}...")
+        
+        # overwrite with new records having inferred category feature
+        with open('/content/drive/MyDrive/ct_data23/processed_trec_topic_X.jsonl', 'w') as f:
+            i = start
+            print(f'starting at: {i}')
+            while i < len(data):
+                next_chunk_end = min(len(data), i+chunk_size)
+                conditions = stream_condition_data(data[i:next_chunk_end], doc_or_topic=doc_or_topic)
+                categories = gen_categories(pipe, conditions)
+                print(f"generated {len(categories)} categories for {chunk_size} conditions...")
+                for j in range(i, next_chunk_end):
+                    data[j][category_label] = categories[j - i]
+                    f.write(json.dumps(data[j]))
+                    f.write('\n')
+                    
+                if doc_or_topic == 'doc':
+                    print(f"{i=}, doc condition: {data[i]['condition']}, generated category: {data[i]['category'].items()}")
+                else:
+                    print(f"{i=}, topic raw text condition: {data[i]['raw_text']}, generated category: {data[i]['category'].items()}")
+                    
+                i += chunk_size
+		
+
+def gen_categories(pipe, text_dataset: Generator[str, None, None]) -> str:
+	categories = []
+	for output in pipe(text_dataset, candidate_labels=CT_CATEGORIES, batch_size=64):
+		score_dict = {output['labels'][i]:output['scores'][i] for i in range(len(output['labels']))}
+		#category = max(score_dict, key=score_dict.get)
+		categories.append(score_dict)
+	return categories
+
+
+def gen_single_category_vector(pipe, text: str) -> str:
+    output = pipe(text, candidate_labels=CT_CATEGORIES)
+    score_dict = {output['labels'][i]:output['scores'][i] for i in range(len(output['labels']))}
+    return np.array(sorted(score_dict, key=score_dict.get, reverse=True))

ctmatch/scripts/get_web_data.py

@@ -0,0 +1,14 @@
+
+from selenium import webdriver
+
+def save_web_data(url: str) -> None:
+	driver = webdriver.Chrome()
+	driver.get(url)
+	button = driver.find_element_by_class_name("save-list")
+	button.click()
+
+
+if __name__ == "__main__":
+	url = "https://clinicaltrials.gov/ct2/results?cond=Heart+Diseases"
+	save_web_data(url)
+

ctmatch/scripts/split_files.py

@@ -0,0 +1,37 @@
+
+
+from pathlib import Path
+import argparse
+import os
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('folder', 
+                    help="supply a folder path to be split up. if not folder, method won't do anything")
+
+args = parser.parse_args()
+
+MAX_FOLDER_SIZE = 2000
+
+
+def split_files(folder: Path):
+
+    assert folder.is_dir()
+    num_dirs = 1
+    curr_size = 0
+
+    new_subfolder_path = folder.parent / f"{folder.as_posix()}_{num_dirs}"
+    new_subfolder_path.mkdir(exist_ok=True)
+    for file in folder.iterdir():
+        if curr_size > MAX_FOLDER_SIZE:
+            num_dirs += 1
+            new_subfolder_path = folder.parent / f"{folder.as_posix()}_{num_dirs}"
+            new_subfolder_path.mkdir(exist_ok=True)
+            curr_size = 0
+        else:
+            curr_size += 1
+        file.rename(new_subfolder_path / file.name)
+
+if __name__ == "__main__":
+    split_files(Path(args.folder))
+

ctmatch/scripts/vis_script.py

@@ -0,0 +1,334 @@
+
+from typing import Dict, List, NamedTuple
+from ctproc.proc import CTDocument, EligCrit
+from matplotlib import pyplot as plt
+from collections import defaultdict
+from zipfile import ZipFile
+from lxml import etree
+import pandas
+import re
+
+
+from utils.ctmatch_utils import *
+
+class FieldCounter(NamedTuple):
+  missfld_counts: Dict[str, int] = defaultdict(int)
+  emptfld_counts: Dict[str, int] = defaultdict(int)
+  elig_form_counts: Dict[str, int] = defaultdict(int)
+  unit_counts: Dict[str, int] = defaultdict(int)
+
+
+#----------------------------------------------------------------#
+# EDA Utility Functions
+#----------------------------------------------------------------#
+
+# viewing
+def print_elig_result(doc, dont_print=[]):
+  for k, v in doc.elig_crit.__dict__.items():
+    if k in dont_print:
+      continue
+    if type(v) == list:
+      print('\n' + k)
+      for v_i in v:
+        print(v_i)
+    else:
+      print(f"{k}: {v}")
+
+
+def display_elig(docs: List[CTDocument]) -> None:
+  age_in_elig_text_dist = count_elig_crit_age_in_text(docs)
+  total = sum(age_in_elig_text_dist.values())
+  print(f"{total} out of {len(docs)} documents had age in eligibility text: {total / len(docs)}%")
+
+  age_in_elig_counts_df = pandas.DataFrame(age_in_elig_text_dist, index=[0])
+  age_in_elig_counts_df.plot(kind="bar", xticks=[], xlabel="include_or_exclude", ylabel="count", title="Age in Eligibility Criteria Text Distribution")
+  print(age_in_elig_counts_df)
+  inc_ratio = age_in_elig_text_dist['inc_ct'] / total
+  exc_ratio = age_in_elig_text_dist['exc_ct'] / total
+  print(f"{age_in_elig_text_dist['inc_ct']} instances in inclusion statements ({inc_ratio}%), {age_in_elig_text_dist['exc_ct']} instances in exclusion statements ({exc_ratio}%)")
+
+
+
+def get_lengths(processed_docs: List[Dict[str, str]]) -> None:
+    no_crit, miss_inc, miss_exc = 0, 0, 0
+    inc_lens, exc_lens, all_lens = 0, 0, 0
+    
+    for i, d in enumerate(processed_docs):
+        crit = d['elig_crit']['raw_text']
+        inc_crit = d['elig_crit']['include_criteria']
+        exc_crit = d['elig_crit']['exclude_criteria']
+        
+        if len(inc_crit) == 0:
+            miss_inc += 1
+
+        if len(exc_crit) == 0:
+            miss_exc += 1
+            
+        if (len(exc_crit) == 0) and (len(inc_crit) == 0):
+            no_crit += 1
+        
+        #print(crit)
+
+        inc_length = sum([len(c.split()) for c in inc_crit])
+        exc_length = sum([len(c.split()) for c in exc_crit])
+        crit_len = inc_length + exc_length
+        inc_lens += inc_length
+        exc_lens += exc_length
+        all_lens += crit_len
+        
+        print(f"{miss_inc=}, {miss_exc=}, {no_crit=}, {inc_lens / len(processed_docs)}, {exc_lens / len(processed_docs)}, {all_lens / len(processed_docs)}")
+
+        
+
+def print_ent_sent(ent_sent):
+  for e in ent_sent:
+    e_small = {}
+    e_small['raw_text'] = e['raw_text']
+    e_small['start'] = e['start']
+    e_small['end'] = e['end']
+    e_small['negation'] = e['negation']
+    print(e_small.items())
+
+
+
+
+
+
+#--------------------------------------------------------------------------------------#
+# methods for getting counts
+#--------------------------------------------------------------------------------------#
+
+def process_counts(zip_data: str) -> FieldCounter:
+  """
+  desc:       main method for processing a zipped file of clinical trial XML documents from clinicaltrials.gov
+              parameterized by CTConfig the self ClinProc object was initialized with
+  returns:    yields processed CTDocuments one at a time
+  """
+
+  counts = FieldCounter()
+  with ZipFile(zip_data, 'r') as zip_reader:
+    for i, ct_file in enumerate(zip_reader.namelist()):
+      if i % 1000 == 0:
+        print(f"{i} docs processed")
+
+      if not ct_file.endswith('xml'):
+        continue
+
+      counts = get_ct_file_counts(zip_reader.open(ct_file), counts)
+  return counts
+
+
+
+
+def get_ct_file_counts(xml_filereader, counts: FieldCounter) -> FieldCounter:
+  doc_tree = etree.parse(xml_filereader)
+  root = doc_tree.getroot()
+
+  # adding new keys vs subdictionaries?????
+  required_fields = {
+    "id":None, 
+    "brief_title":None, 
+    "eligibility/criteria/textblock":None, 
+    "eligibility/gender":"Default Value", 
+    "eligibility/minimum_age":{"male":0, "female":0}, 
+    "eligibility/maximum_age":{"male":999., "female":999.}, 
+    "detailed_description/textblock":None, 
+    "condition":None,
+    "condition/condition_browse":None,
+    "intervention/intervention_type":None,
+    "intervention/intervention_name":None,
+    "intervention_browse/mesh_term":None,
+    "brief_summary/textblock":None,
+  }
+        
+  for field in required_fields.keys():
+    field_tag = 'id_info/nct_id'  if field == 'id' else field
+    try:
+      field_val = root.find(field_tag).text
+      if not EMPTY_PATTERN.fullmatch(field_val):
+        if field == 'eligibility/criteria/textblock':
+          counts.elig_form_counts = get_elig_counts(field_val, counts.elig_form_counts)
+        elif "age" in field:
+          age_match = AGE_PATTERN.match(field_val)
+          if age_match is not None:
+            unit = age_match.group('units')
+            if unit is not None:
+              counts.unit_counts[unit] += 1
+
+
+
+    except:
+      if root.find(field_tag) is None:
+        counts.missfld_counts[field]  += 1
+      elif EMPTY_PATTERN.fullmatch(root.find(field_tag).text):
+        counts.emptfld_counts[field]  += 1
+    
+  return counts
+
+
+
+
+
+
+
+
+
+
+
+
+def get_elig_counts(elig_text: str, elig_form_counts: Dict[str, int]) -> Dict[str, int]:
+  assert elig_text is not None, "Eligibility text is empty"
+  if re.search('[Ii]nclusion [Cc]riteria:[^\w]+\n', elig_text):
+    if re.search('[Ee]xclusion Criteria:[^\w]+\n', elig_text):
+      elig_form_counts["inc_and_exc"] += 1
+      return elig_form_counts
+    else:
+      elig_form_counts["inc_only"] += 1
+      return elig_form_counts
+
+  elif re.search('[Ee]xclusion [Cc]riteria:[^\w]+\n', elig_text):
+    elig_form_counts["exc_only"] += 1
+    return elig_form_counts
+  
+  else:
+    elig_form_counts["textblock"] += 1
+    return  elig_form_counts
+
+
+
+
+def get_counts(docs: List[CTDocument]):
+  gender_dist = defaultdict(int)
+  min_age_dist = defaultdict(int)
+  max_age_dist = defaultdict(int)
+  for doc in docs:
+    gender_dist[doc.elig_gender] += 1
+    min_age_dist[doc.elig_crit.elig_min_age] += 1
+    max_age_dist[doc.elig_max_age] += 1
+  return gender_dist, min_age_dist, max_age_dist
+  
+
+
+def get_relled(topic_id, rel_dict):
+  twos, ones, zeros = set(), set(), set()
+  for doc_id, rel in rel_dict[topic_id].items():
+    if rel == 1:
+      ones.add(doc_id)
+    elif rel == 2:
+      twos.add(doc_id)
+    else:
+      zeros.add(doc_id)
+  return {"twos": twos, "ones": ones, "zeros": zeros}
+  
+def scan_for_age(
+  elig_crit: EligCrit, 
+  inc_or_ex: str = 'include'
+) -> bool:
+  crit_to_scan = elig_crit.include_criteria if inc_or_ex == 'include' else elig_crit.exclude_criteria
+  for crit in crit_to_scan:
+    if re.match(r' ages? ', crit.lower()) is not None:
+      return True
+  return False
+
+
+def count_elig_crit_age_in_text(docs, skip_predefined:bool = True):
+  age_in_elig_text_dist = defaultdict(int)
+  skipped = 0
+  for doc in docs:
+    if skip_predefined:
+      if (doc.elig_min_age != 0) or (doc.elig_max_age != 999):  # author(s) have specified SOME criteria, assumes judgment prefers this field to free trex in criteria textblock
+        skipped += 1
+        continue
+
+    age_in_elig_text_dist['include'] += scan_for_age(doc.elig_crit, 'include')
+    age_in_elig_text_dist['exclude'] += scan_for_age(doc.elig_crit, 'exclude') 
+    
+  print(f"Total skipped: {skipped}")
+  return age_in_elig_text_dist
+
+
+
+
+def get_missing_criteria(docs: List[CTDocument]):
+  missing_inc_ids, missing_exc_ids = {}, {}
+  for d in docs:
+
+    if len(d.elig_crit.include_criteria) == 0:
+      missing_inc_ids.add(d.nct_id)
+  
+    if len(d.elig_crit.exclude_criteria) == 0:
+      missing_exc_ids.add(d.nct_id)
+      
+  return missing_inc_ids, missing_exc_ids
+
+
+# for evaluating effect of filtering
+def get_doc_percent_elig(filtered_docs_by_topic: Dict[str, set]):
+  percents_elig = []
+  for topic_id, doc_list in filtered_docs_by_topic.items():
+    per = len(doc_list) / 3262.0
+    percents_elig.append(per)
+    print(topic_id, len(doc_list), per)
+  mean_elig = sum(percents_elig) / len(percents_elig)
+  print(f"Mean elgibile number of docs: {mean_elig}")
+
+
+
+
+
+# plotting
+
+def plot_counts(missfld_counts, emptfld_counts):
+  miss_df = pandas.DataFrame(missfld_counts, index=[0])
+  miss_df.plot(kind='bar', xticks=[], title="Missing Fields", ylabel="count", xlabel="field")
+  plt.legend(loc=(1.04, 0))
+
+  empt_df = pandas.DataFrame(emptfld_counts, index=[0])
+  empt_df.plot(kind='bar', xticks=[], title="Empty Fields", ylabel="count", xlabel="field")
+  plt.legend(loc=(1.04, 0))
+
+
+
+
+#----------------------------------------------------------------#
+# EDA Test Data Utility Functions
+#----------------------------------------------------------------#
+
+
+def get_test_rels(test_rels):
+    rel_dict = defaultdict(lambda:defaultdict(int))
+    rel_type_dict = defaultdict(int)
+    for line in open(test_rels, 'r').readlines():
+        topic_id, _, doc_id, rel = re.split(r'\s+', line.strip())
+        rel_dict[topic_id][doc_id] = int(rel)
+        rel_type_dict[rel] += 1
+    return rel_dict, rel_type_dict
+
+def analyze_test_rels(test_rels_path):
+    rel_dict, rel_type_dict = get_test_rels(test_rels_path)
+
+    print("Rel Type Results:")
+    for t, n in rel_type_dict.items():
+      print(t + ': ' + str(n))
+
+    lengths = dict()
+    all_qrelled_docs = set()
+    for tid in rel_dict.keys():
+        lengths[tid] = len(rel_dict[tid])
+        for d in rel_dict[tid].keys():
+            all_qrelled_docs.add(d)
+    for topic, num_relled in lengths.items():
+        print(topic, num_relled)
+    print(f"Total relled: {len(all_qrelled_docs)}")
+    return rel_type_dict, rel_dict, all_qrelled_docs
+
+
+
+
+
+
+if __name__ == '__main__':
+	qrels_path = '/Users/jameskelly/Documents/cp/ctmatch/data/qrels-clinical_trials.txt'
+	rel_type_dict, rel_dict, all_qrelled_docs = analyze_test_rels(qrels_path)
+  #docs_path = '/Users/jameskelly/Documents/cp/ctproc/clinicaltrials.gov-16_dec_2015_17.zip'
+  #counts = process_counts(docs_path)

ctmatch/utils/ctmatch_utils.py

@@ -0,0 +1,133 @@
+
+from typing import Any, Dict, List, Optional, Set
+from sklearn.metrics.pairwise import linear_kernel
+from collections import defaultdict
+from numpy.linalg import norm
+from datasets import Dataset
+import numpy as np
+import json
+import re
+
+
+
+
+
+#----------------------------------------------------------------#
+# global regex patterns for use throughout the methods
+#----------------------------------------------------------------#
+
+
+EMPTY_PATTERN = re.compile('[\n\s]+')
+"""
+both_inc_and_exc_pattern = re.compile(r\"\"\"[\s\n]*[Ii]nclusion [Cc]riteria:?               # top line of both
+                                      (?:[ ]+[Ee]ligibility[ \w]+\:[ ])?                  # could contain this unneeded bit next
+                                      (?P<include_crit>[ \n\-\.\?\"\%\r\w\:\,\(\)]*)      # this should get all inclusion criteria as a string
+                                      [Ee]xclusion[ ][Cc]riteria:?                        # delineator to exclusion criteria
+                                      (?P<exclude_crit>[\w\W ]*)                          # exclusion criteria as string
+                                      \"\"\", re.VERBOSE)
+"""
+INC_ONLY_PATTERN = re.compile('[\s\n]+[Ii]nclusion [Cc]riteria:?([\w\W ]*)')
+EXC_ONLY_PATTERN = re.compile('[\n\r ]+[Ee]xclusion [Cc]riteria:?([\w\W ]*)')
+AGE_PATTERN = re.compile('(?P<age>\d+) *(?P<units>\w+).*')
+YEAR_PATTERN = re.compile('(?P<year>[yY]ears?.*)')
+MONTH_PATTERN = re.compile('(?P<month>[mM]o(?:nth)?)')
+WEEK_PATTERN = re.compile('(?P<week>[wW]eeks?)')
+
+BOTH_INC_AND_EXC_PATTERN = re.compile("[\s\n]*[Ii]nclusion [Cc]riteria:?(?: +[Ee]ligibility[ \w]+\: )?(?P<include_crit>[ \n\-\.\?\"\%\r\w\:\,\(\)]*)[Ee]xclusion [Cc]riteria:?(?P<exclude_crit>[\w\W ]*)")
+
+
+
+# -------------------------------------------------------------------------------------- #
+# pretokenization utils (should be in a tokenizer...)
+# -------------------------------------------------------------------------------------- #
+
+def truncate(s: str, max_tokens: Optional[int] = None) -> str:
+  if max_tokens is None:
+    return s
+  s_tokens = s.split()
+  return ' '.join(s_tokens[:min(len(s_tokens), max_tokens)])
+
+
+
+# -------------------------------------------------------------------------------------- #
+# I/O utils
+# -------------------------------------------------------------------------------------- #
+
+def save_docs_jsonl(docs: List[Any], writefile: str) -> None:
+  """
+  desc:    iteratively writes contents of docs as jsonl to writefile 
+  """
+  with open(writefile, "w") as outfile:
+    for doc in docs:
+      json.dump(doc, outfile)
+      outfile.write("\n")
+
+
+def get_processed_data(proc_loc: str, get_only: Optional[Set[str]] = None):
+  """
+  proc_loc:    str or path to location of docs in jsonl form
+  """
+  with open(proc_loc, 'r') as json_file:
+    json_list = list(json_file)
+
+  if get_only is None:
+    for json_str in json_list:
+      yield json.loads(json_str)
+
+  else:
+    for s in json_list:
+      s_data = json.loads(s)
+      if s_data["id"] in get_only:
+        yield s_data
+        get_only.remove(s_data['id'])
+        if len(get_only) == 0:
+          return
+
+  
+  
+
+
+def train_test_val_split(dataset, splits: Dict[str, float], seed: int = 37) -> Dataset:
+  """
+  splits a dataset having only "train" into one having train, test, val, with 
+  split sizes determined by splits["train"] and splits["val"] (dict must have those keys)
+
+  """
+  dataset = dataset["train"].train_test_split(train_size=splits["train"], seed=seed)
+  train = dataset["train"]
+  sub = train.train_test_split(test_size=splits["val"],  seed=seed)
+  new_train = sub["train"]
+  new_val = sub["test"]
+  dataset["train"] = new_train
+  dataset["validation"] = new_val
+  return dataset
+
+
+
+#----------------------------------------------------------------#
+# computation methods
+#----------------------------------------------------------------#
+
+def exclusive_argmax(vector: np.ndarray) -> np.ndarray:
+    mask = np.zeros(len(vector))
+    argmax = np.argmax(vector)
+    vector = vector * mask
+    vector[argmax]  = 1
+    return vector    
+
+
+#----------------------------------------------------------------#
+# evaluation methods (duplicated from ctproc scripts)
+#----------------------------------------------------------------#
+
+def get_test_rels(rel_path):
+    rel_dict = defaultdict(lambda:defaultdict(int))
+    rel_type_dict = defaultdict(int)
+    for line in open(rel_path, 'r').readlines():
+        topic_id, _, doc_id, rel = re.split(r'\s+', line.strip())
+        rel_dict[topic_id][doc_id] = int(rel)
+        rel_type_dict[rel] += 1
+    return rel_dict, rel_type_dict
+
+
+

ctmatch/utils/eval_utils.py

@@ -0,0 +1,91 @@
+
+from typing import Dict, List, Tuple
+
+from sklearn.metrics import f1_score
+from collections import defaultdict
+from lxml import etree
+import numpy as np
+
+  
+def get_trec_topic2text(topic_path) -> Dict[str, str]:
+    """
+    desc:       main method for processing a single XML file of TREC21 patient descriptions called "topics" in this sense
+    returns:    dict of topicid: topic text
+    """
+
+    topic2text = {}
+    topic_root = etree.parse(topic_path).getroot()
+    for topic in topic_root:
+        topic2text[topic.attrib['number']] = topic.text
+    
+    return topic2text
+    
+
+
+def get_kz_topic2text(topic_path) -> Dict[str, str]:
+    """
+    desc:       main method for processing a single XML file of TREC21 patient descriptions called "topics" in this sense
+    returns:    dict of topicid: topic text
+    """
+
+    topic2text = {}
+    with open(topic_path, 'r') as f:
+        for line in f.readlines():
+            line = line.strip()
+            
+            if line.startswith('<TOP>'):
+                topic_id, text = None, None
+                continue
+
+            if line.startswith('<NUM>'):
+                topic_id = line[5:-6]
+
+            elif line.startswith('<TITLE>'):
+                text = line[7:].strip()
+                topic2text[topic_id] = text
+
+    return topic2text
+          
+
+
+def calc_first_positive_rank(ranked_ids: List[str], doc2rel: Dict[str, int], pos_val: int = 2) -> Tuple[int, float]:
+    """
+    desc:       compute the mean reciprocal rank of a ranking
+    returns:    mrr 
+    """
+    for i, doc_id in enumerate(ranked_ids):
+        if doc2rel[doc_id] == pos_val:
+            return i + 1, 1./float(i+1)
+    return len(ranked_ids) + 1, 0.0
+
+
+def calc_f1(ranked_ids: List[str], doc2rel: Dict[str, int]) -> Dict[str, Dict[str, float]]:
+  label_counts = get_label_counts(doc2rel)
+  predicted, ground_truth = [], []
+  for doc_id in ranked_ids:
+    # 2, 1, 0
+    ground_truth.append(doc2rel[doc_id])
+    pred_label = get_predicted_label(label_counts)
+    predicted.append(pred_label)
+    label_counts[pred_label] -= 1
+
+  return f1_score(ground_truth, predicted, average='micro')
+
+  
+
+def get_label_counts(doc2rel: Dict[str, int]) -> Dict[int, int]:
+  """
+  return an ordered list of [(2, <count_2s>), (1, <count_1s>), (0, count_0s)]
+  """
+  label_counts = defaultdict(int)
+  for scored_doc in doc2rel:
+    label = doc2rel[scored_doc]
+    label_counts[label] += 1
+  return label_counts
+
+def get_predicted_label(label_counts: Dict[int, int]) -> int:
+  if label_counts[2] > 0:
+    return 2
+  if label_counts[1] > 0:
+    return 1
+  return 0

requirements.txt

@@ -0,0 +1,23 @@
+#ctproc  uncoment if doing data prep on raw ct documents
+#https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.4.0/en_core_sci_md-0.4.0.tar.gz  uncomment if using ctproc
+#pyserini==0.12.0       uncomment if using ctproc with indexes (not recommended)
+#git+https://github.com/semajyllek/transformers.git@add-biogpt-sequenceclassifier
+#sacremoses uncomment if using biogpt
+sentence-transformers
+huggingface_hub
+scikit-learn
+transformers
+onnxruntime
+nn_pruning
+optimum
+onnx
+
+matplotlib
+accelerate
+datasets
+evaluate
+pandas
+openai
+lxml
+
+gradio