Claim2Source/startup.py at master · faerber-lab/Claim2Source · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
from retrieval.bm25 import BM25RetrievalSystem
from retrieval.sentence_embeddings import SentenceTransformerRetrievalSystem
from retrieval.gritlm import GritLMRetrievalSystem
from ranking.ranking import ranking
from evaluation.evaluation import evaluate_retrieval_systems, checkthat_submission


def main():
    k_queries = -1  # Specify number of queries to use (-1: all)
    synthetic_abstracts = False  # Set to True if style transfer should be used (Specify style in config.py)
    synthetic_queries = False  # Set to True if style transfer should be used (Specify style in config.py)

    # List of all retrieval systems for experiments
    retrieval_systems = [BM25RetrievalSystem(),
                         SentenceTransformerRetrievalSystem(model_name="all-MiniLM-L6-v2"),
                         SentenceTransformerRetrievalSystem(model_name="all-mpnet-base-v2"),
                         SentenceTransformerRetrievalSystem(model_name="intfloat/e5-large-v2"),
                         SentenceTransformerRetrievalSystem(model_name="gtr-t5-xl"),
                         GritLMRetrievalSystem(),
                         SentenceTransformerRetrievalSystem(model_name="malteos/scincl"),
                         SentenceTransformerRetrievalSystem(model_name="allenai/specter")]

    try:
        # Create Faiss index
        st = SentenceTransformerRetrievalSystem(model_name="intfloat/e5-large-v2")
        st.save_corpus_embeddings(synthetic_abstracts=synthetic_abstracts)
        st = SentenceTransformerRetrievalSystem(model_name="all-MiniLM-L6-v2")
        st.save_corpus_embeddings(synthetic_abstracts=synthetic_abstracts)
        st = SentenceTransformerRetrievalSystem(model_name="all-mpnet-base-v2")
        st.save_corpus_embeddings(synthetic_abstracts=synthetic_abstracts)
        st = SentenceTransformerRetrievalSystem(model_name="gtr-t5-xl")
        st.save_corpus_embeddings(synthetic_abstracts=synthetic_abstracts)
        st = GritLMRetrievalSystem()
        st.save_corpus_embeddings(synthetic_abstracts=synthetic_abstracts)
        st = SentenceTransformerRetrievalSystem(model_name="malteos/scincl")
        st.save_corpus_embeddings(synthetic_abstracts=synthetic_abstracts)
        st = SentenceTransformerRetrievalSystem(model_name="allenai/specter")
        st.save_corpus_embeddings(synthetic_abstracts=synthetic_abstracts)

        # Run ranking
        ranking(retrieval_system=BM25RetrievalSystem(), k_queries=k_queries, synthetic_queries=synthetic_queries,
                synthetic_abstracts=synthetic_abstracts)
        ranking(retrieval_system=SentenceTransformerRetrievalSystem(model_name="intfloat/e5-large-v2"), k_queries=k_queries,
                synthetic_queries=synthetic_queries, synthetic_abstracts=synthetic_abstracts)
        ranking(retrieval_system=SentenceTransformerRetrievalSystem(model_name="all-MiniLM-L6-v2"), k_queries=k_queries,
                synthetic_queries=synthetic_queries, synthetic_abstracts=synthetic_abstracts)
        ranking(retrieval_system=SentenceTransformerRetrievalSystem(model_name="all-mpnet-base-v2"), k_queries=k_queries,
                synthetic_queries=synthetic_queries, synthetic_abstracts=synthetic_abstracts)
        ranking(retrieval_system=SentenceTransformerRetrievalSystem(model_name="gtr-t5-xl"), k_queries=k_queries,
                synthetic_queries=synthetic_queries, synthetic_abstracts=synthetic_abstracts)
        ranking(retrieval_system=GritLMRetrievalSystem(), k_queries=k_queries,
                synthetic_queries=synthetic_queries, synthetic_abstracts=synthetic_abstracts)
        ranking(retrieval_system=SentenceTransformerRetrievalSystem(model_name="malteos/scincl"), k_queries=k_queries,
                synthetic_queries=synthetic_queries, synthetic_abstracts=synthetic_abstracts)
        ranking(retrieval_system=SentenceTransformerRetrievalSystem(model_name="allenai/specter"), k_queries=k_queries,
                synthetic_queries=synthetic_queries, synthetic_abstracts=synthetic_abstracts)

        # Evaluate retrieval
        evaluate_retrieval_systems(retrieval_systems=retrieval_systems, k_queries=-1, synthetic_queries=True,
                                   synthetic_abstracts=False)

    except Exception as e:
        print(e)

    pass


if __name__ == '__main__':
    main()