init

README.md

@@ -0,0 +1,12 @@
+This repo is modified from [Quin+](https://github.com/algoprog/Quin).
+
+To set up:
+1. Copy `../models` folder into this folder
+2. Follow instructions in the original repository to download the models.
+3. Install requirements. The requirements.txt file in the original repo were for Python 3.7, but the requirements in this repo are based on Python 3.8. Choose according to the Python version available on your system.
+4. Run
+```shell
+python quin_web.py
+```
+
+To check if the backend is running correctly, set up frontend and enter a query there. By default at `INFO` lever, the logger will record each incoming query.

app.py

@@ -0,0 +1,277 @@
+import argparse
+import gradio as gr
+import json
+import logging
+import nltk
+import numpy
+import os
+import re
+import requests
+import torch
+import urllib
+import urllib.parse as urlparse
+
+from bs4 import BeautifulSoup
+from flask import request, Flask, jsonify
+from flask_cors import CORS
+from huggingface_hub import hf_hub_download, snapshot_download, login
+from models.nli import NLI
+from models.qa_ranker import PassageRanker
+from models.text_encoder import SentenceTransformer
+from multiprocessing.pool import ThreadPool
+from scipy import spatial
+from scipy.special import softmax
+from time import sleep
+from urllib.parse import parse_qs
+from web_search import get_html, WebParser, duckduckgo_search
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+API_URLS = ['https://letscheck.nus.edu.sg/quin/']
+logging.getLogger().setLevel(logging.INFO)
+
+
+def is_question(query):
+    if re.match(r'^(who|when|what|why|which|whose|is|are|was|were|do|does|did|how)', query) or query.endswith('?'):
+        return True
+    pos_tags = nltk.pos_tag(nltk.word_tokenize(query))
+    for tag in pos_tags:
+        if tag[1].startswith('VB'):
+            return False
+    return True
+
+
+def clear_cache():
+    for url in API_URLS:
+        print(f"Clearing cache from {url}")
+        r = requests.post(f'{url}/clear_cache?secret=123')
+        print(r.status_code)
+        print(r)
+
+
+class Quin:
+    def __init__(self):
+        nltk.download('punkt')
+        self.sent_tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
+        self.sent_tokenizer._params.abbrev_types.update(['e.g', 'i.e', 'subsp'])
+        self.app = Flask(__name__)
+        CORS(self.app)
+        
+        snapshot_download(repo_id="anabmaulana/qrbert-multitask", local_dir="cache/qrbert-multitask")
+        hf_hub_download(repo_id="anabmaulana/quin-passage-ranker", filename="passage_ranker.state_dict", local_dir="cache/quin-passage-ranker")
+        hf_hub_download(repo_id="anabmaulana/quin-nli", filename="nli.state_dict", local_dir="cache/quin-nli")
+
+        # torch.cuda.is_available = lambda: False
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+        self.text_embedding_model = SentenceTransformer('cache/qrbert-multitask',
+                                                        device=device,
+                                                        parallel=True)
+        self.passage_ranking_model = PassageRanker(
+            model_path='cache/quin-passage-ranker/passage_ranker.state_dict',
+            gpu=torch.cuda.is_available(),
+            batch_size=16)
+        self.passage_ranking_model.eval()
+        self.nli_model = NLI('cache/quin-nli/nli.state_dict',
+                             batch_size=64,
+                             device=device,
+                             parallel=True)
+        self.nli_model.eval()
+        logging.info('Initialized!')
+
+    def extract_snippets(self, text, sentences_per_snippet=4):
+        sentences = self.sent_tokenizer.tokenize(text)
+        snippets = []
+        i = 0
+        last_index = 0
+        while i < len(sentences):
+            snippet = ' '.join(sentences[i:i + sentences_per_snippet])
+            if len(snippet.split(' ')) > 4:
+                snippets.append(snippet)
+            last_index = i + sentences_per_snippet
+            i += sentences_per_snippet
+        if last_index < len(sentences):
+            snippet = ' '.join(sentences[last_index:])
+            if len(snippet.split(' ')) > 4:
+                snippets.append(snippet)
+        return snippets
+
+    def search_web_evidence(self, query, limit=10):
+        logging.info('searching the web...')
+        urls = duckduckgo_search(query, pages=1)[:limit]
+        logging.info('downloading {} web pages...'.format(len(urls)))
+        search_results = []
+        # Move this out of the function, since it is used a few times
+        query_is_question = is_question(query)
+
+        # Local methods for multithreading
+        def download(url):
+            nonlocal search_results
+            data = get_html(url)
+            soup = BeautifulSoup(data, features='lxml')
+            title = soup.title.string
+            w = WebParser()
+            w.feed(data)
+            new_snippets = sum([self.extract_snippets(b) for b in w.blocks if b.count(' ') > 20], [])
+            new_snippets = [{'snippet': snippet, 'url': url, 'title': title} for snippet in new_snippets]
+            search_results += new_snippets
+
+        def timeout_download(arg):
+            pool = ThreadPool(1)
+            try:
+                pool.apply_async(download, [arg]).get(timeout=2)
+            except:
+                pass
+            pool.close()
+            pool.join()
+
+        p = ThreadPool(32)
+        p.map(timeout_download, urls)
+        p.close()
+        p.join()
+
+        if query_is_question:
+            logging.info('re-ranking...')
+            snippets = [s['snippet'] for s in search_results]
+            qa_pairs = [(query, snippet) for snippet in snippets]
+            _, probs = self.passage_ranking_model(qa_pairs)
+            probs = [softmax(p)[1] for p in probs]
+            filtered_results = []
+            for i in range(len(search_results)):
+                if probs[i] > 0.35:
+                    search_results[i]['score'] = str(probs[i])
+                    filtered_results.append(search_results[i])
+            search_results = filtered_results
+            search_results = sorted(search_results, key=lambda x: float(x['score']), reverse=True)
+
+        # Split into a different function below
+        # highlight_results(self, search_results)
+        # highlight most relevant sentences
+        logging.info('highlighting...')
+        results_sentences = []
+        sentences_texts = []
+        sentences_vectors = {}
+        for i, r in enumerate(search_results):
+            sentences = self.sent_tokenizer.tokenize(r['snippet'])
+            sentences = [s for s in sentences if len(s.split(' ')) > 4]
+            sentences_texts.extend(sentences)
+            results_sentences.append(sentences)
+
+        vectors = self.text_embedding_model.encode(sentences=sentences_texts, batch_size=128)
+        for i, v in enumerate(vectors):
+            sentences_vectors[sentences_texts[i]] = v
+
+        query_vector = self.text_embedding_model.encode(sentences=[query], batch_size=1)[0]
+        for i, sentences in enumerate(results_sentences):
+            best_sentences = set()
+            evidence_sentences = []
+            for sentence in sentences:
+                sentence_vector = sentences_vectors[sentence]
+                score = 1 - spatial.distance.cosine(query_vector, sentence_vector)
+                if score > 0.91:
+                    best_sentences.add(sentence)
+                    evidence_sentences.append(sentence)
+            if len(evidence_sentences) > 0:
+                search_results[i]['evidence'] = ' '.join(evidence_sentences)
+            search_results[i]['snippet'] = \
+                ' '.join([s if s not in best_sentences else '<b>{}</b>'.format(s) for s in sentences])
+
+        search_results = [s for s in search_results if 'evidence' in s]
+
+        # Split into a different function below
+        # entailment_classification(self, search_results)
+        if not query_is_question:
+            logging.info('entailment classification...')
+            es_pairs = []
+            for result in search_results:
+                evidence = result['evidence']
+                es_pairs.append((evidence, query))
+            try:
+                labels, probs = self.nli_model(es_pairs)
+                logging.info(str(labels))
+                for i in range(len(labels)):
+                    confidence = numpy.exp(numpy.max(probs[i]))
+                    if confidence > 0.4:
+                        search_results[i]['nli_class'] = labels[i]
+                    else:
+                        search_results[i]['nli_class'] = 'neutral'
+                    search_results[i]['nli_confidence'] = str(confidence)
+            except Exception as e:
+                logging.warning('Error doing entailment classification')
+                logging.warning(str(e))
+            try:
+                search_results = [s for s in search_results if s['nli_class'] != 'neutral']
+            except:
+                pass
+
+        filtered_results = []
+        added_urls = set()
+        supporting = 0
+        refuting = 0
+        for r in search_results:
+            if r['url'] not in added_urls:
+                filtered_results.append(r)
+                added_urls.add(r['url'])
+                if 'nli_class' in r:
+                    if r['nli_class'] == 'entailment':
+                        supporting += 1
+                    elif r['nli_class'] == 'contradiction':
+                        refuting += 1
+        search_results = filtered_results
+
+        if supporting > 4 or refuting > 4:
+            if supporting / (refuting + 0.001) > 1.7:
+                veracity = 'Probably True'
+            elif refuting / (supporting + 0.001) > 1.7:
+                veracity = 'Probably False'
+            else:
+                veracity = '? Ambiguous'
+        else:
+            veracity = 'Not enough evidence'
+
+        search_results = search_results[:limit]
+
+        logging.info('done searching')
+
+        if query_is_question:
+            return {'type': 'question',
+                    'results': search_results}
+        else:
+            return {'type': 'statement',
+                    'supporting': supporting,
+                    'refuting': refuting,
+                    'veracity_rating': veracity,
+                    'results': search_results}
+
+    def build_endpoints(self):
+        @self.app.route('/search', methods=['POST', 'GET'])
+        def search_endpoint():
+            query = request.args.get('query').lower()
+            limit = request.args.get('limit') or 100
+            limit = min(int(limit), 100)
+            results = json.dumps(self.search_web_evidence(query, limit=limit), indent=4)
+            return results
+
+    def serve(self, port=12345):
+        self.build_endpoints()
+        self.app.run(host='0.0.0.0', port=port)
+
+
+
+fchecker = Quin()
+
+def predict(query, limit):
+    return fchecker.search_web_evidence(query, limit=limit)
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Textbox(label="Enter Query"),   # string input
+        gr.Number(label="Enter Limit")   # number input
+    ],
+    outputs=gr.JSON(label="Output")
+)
+
+
+if __name__ == '__main__':
+    demo.launch()

models/data_utils.py

@@ -0,0 +1,1006 @@
+import csv
+import json
+import pickle
+
+import logging
+import re
+from pprint import pprint
+
+import pandas
+import gzip
+import os
+
+import numpy as np
+
+from random import randint, random
+from tqdm import tqdm
+from typing import List, Dict
+
+from spacy.lang.en import English
+# from tensorflow.keras.utils import Sequence
+from transformers import RobertaTokenizer, BertTokenizer, BertModel, AutoTokenizer, AutoModel
+
+# from utils import WebParser
+from models.dense_retriever import DenseRetriever
+from models.tokenization import tokenize
+
+from typing import Union, List
+
+from unidecode import unidecode
+
+import spacy
+
+
+# from utils import WebParser
+
+
+class InputExample:
+    """
+    Structure for one input example with texts, the label and a unique id
+    """
+
+    def __init__(self, guid: str, texts: List[str], label: Union[int, float]):
+        """
+        Creates one InputExample with the given texts, guid and label
+        str.strip() is called on both texts.
+        :param guid
+            id for the example
+        :param texts
+            the texts for the example
+        :param label
+            the label for the example
+        """
+        self.guid = guid
+        self.texts = [text.strip() for text in texts]
+        self.label = label
+
+    def get_texts(self):
+        return self.texts
+
+    def get_label(self):
+        return self.label
+
+
+class LoggingHandler(logging.Handler):
+    def __init__(self, level=logging.NOTSET):
+        super().__init__(level)
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            tqdm.write(msg)
+            self.flush()
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except:
+            self.handleError(record)
+
+
+def get_examples(filename, max_examples=0):
+    examples = []
+    id = 0
+    with open(filename, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+            label = sample['label']
+            guid = "%s-%d" % (filename, id)
+            id += 1
+            if label == 'entailment':
+                label = 0
+            elif label == 'contradiction':
+                label = 1
+            else:
+                label = 2
+            examples.append(InputExample(guid=guid,
+                                         texts=[sample['s1'], sample['s2']],
+                                         label=label))
+            if 0 < max_examples <= len(examples):
+                break
+    return examples
+
+
+def get_qa_examples(filename, max_examples=0, dev=False):
+    examples = []
+    id = 0
+    with open(filename, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+            label = sample['relevant']
+            guid = "%s-%d" % (filename, id)
+            id += 1
+            examples.append(InputExample(guid=guid,
+                                         texts=[sample['question'], sample['answer']],
+                                         label=label))
+            if not dev:
+                if label == 1:
+                    for _ in range(13):
+                        examples.append(InputExample(guid=guid,
+                                                     texts=[sample['question'], sample['answer']],
+                                                     label=label))
+            if 0 < max_examples <= len(examples):
+                break
+    return examples
+
+
+def map_label(label):
+    labels = {"relevant": 0, "irrelevant": 1}
+    return labels[label.strip().lower()]
+
+
+def get_qar_examples(filename, max_examples=0):
+    examples = []
+    id = 0
+    with open(filename, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+            guid = "%s-%d" % (filename, id)
+            id += 1
+            examples.append(InputExample(guid=guid,
+                                         texts=[sample['question'], sample['answer']],
+                                         label=1.0))
+            if 0 < max_examples <= len(examples):
+                break
+    return examples
+
+
+def get_qar_artificial_examples():
+    examples = []
+    id = 0
+
+    print('Loading passages...')
+
+    passages = []
+    file = open('data/msmarco/collection.tsv', 'r', encoding='utf8')
+    while True:
+        line = file.readline()
+        if not line:
+            break
+        line = line.rstrip('\n').split('\t')
+        passages.append(line[1])
+
+    print('Loaded passages')
+
+    with open('data/qar/qar_artificial_queries.csv') as f:
+        for i, line in enumerate(f):
+            queries = line.rstrip('\n').split('|')
+            for query in queries:
+                guid = "%s-%d" % ('', id)
+                id += 1
+                examples.append(InputExample(guid=guid,
+                                             texts=[query, passages[i]],
+                                             label=1.0))
+    return examples
+
+
+def get_single_examples(filename, max_examples=0):
+    examples = []
+    id = 0
+    with open(filename, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+            guid = "%s-%d" % (filename, id)
+            id += 1
+            examples.append(InputExample(guid=guid,
+                                         texts=[sample['text']],
+                                         label=1))
+            if 0 < max_examples <= len(examples):
+                break
+    return examples
+
+
+def get_qnli_examples(filename, max_examples=0, no_contradictions=False, fever_only=False):
+    examples = []
+    id = 0
+    with open(filename, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+            label = sample['label']
+            if label == 'contradiction' and no_contradictions:
+                continue
+            if sample['evidence'] == '':
+                continue
+            if fever_only and sample['source'] != 'fever':
+                continue
+            guid = "%s-%d" % (filename, id)
+            id += 1
+
+            examples.append(InputExample(guid=guid,
+                                         texts=[sample['statement'].strip(), sample['evidence'].strip()],
+                                         label=1.0))
+            if 0 < max_examples <= len(examples):
+                break
+    return examples
+
+
+def get_retrieval_examples(filename, negative_corpus='data/msmarco/collection.tsv', max_examples=0, no_statements=True,
+                           encoder_model=None, negative_samples_num=4):
+    examples = []
+    queries = []
+    passages = []
+    negative_passages = []
+    id = 0
+    with open(filename, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+
+            if 'evidence' in sample and sample['evidence'] == '':
+                continue
+
+            guid = "%s-%d" % (filename, id)
+            id += 1
+
+            if sample['type'] == 'question':
+                query = sample['question']
+                passage = sample['answer']
+            else:
+                query = sample['statement']
+                passage = sample['evidence']
+
+            query = query.strip()
+            passage = passage.strip()
+
+            if sample['type'] == 'statement' and no_statements:
+                continue
+
+            queries.append(query)
+            passages.append(passage)
+
+            if sample['source'] == 'natural-questions':
+                negative_passages.append(passage)
+
+            if max_examples == len(passages):
+                break
+
+    if encoder_model is not None:
+        # Load MSMARCO passages
+        logging.info('Loading MSM passages...')
+        with open(negative_corpus) as file:
+            for line in file:
+                p = line.rstrip('\n').split('\t')[1]
+                negative_passages.append(p)
+
+        logging.info('Building ANN index...')
+        dense_retriever = DenseRetriever(model=encoder_model, batch_size=1024, use_gpu=True)
+        dense_retriever.create_index_from_documents(negative_passages)
+        results = dense_retriever.search(queries=queries, limit=100, probes=256)
+        negative_samples = [
+            [negative_passages[p[0]] for p in r if negative_passages[p[0]] != passages[i]][:negative_samples_num]
+            for i, r in enumerate(results)
+        ]
+        # print(queries[0])
+        # print(negative_samples[0][0])
+
+        for i in range(len(queries)):
+            texts = [queries[i], passages[i]] + negative_samples[i]
+            examples.append(InputExample(guid=guid,
+                                         texts=texts,
+                                         label=1.0))
+
+    else:
+        for i in range(len(queries)):
+            texts = [queries[i], passages[i]]
+            examples.append(InputExample(guid=guid,
+                                         texts=texts,
+                                         label=1.0))
+
+    return examples
+
+
+def get_ict_examples(filename, max_examples=0):
+    examples = []
+    id = 0
+    with open(filename, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+            # label = sample['label']
+            guid = "%s-%d" % (filename, id)
+            id += 1
+            examples.append(InputExample(guid=guid,
+                                         texts=[sample['s1'].strip(), sample['s2'].strip()],
+                                         label=1.0))
+            if 0 < max_examples <= len(examples):
+                break
+    return examples
+
+
+def preprocess_fever_dataset(path):
+    def replace_symbols(text):
+        return text.replace('-LRB-', '(') \
+            .replace('-RRB-', ')') \
+            .replace('-LSB-', '[') \
+            .replace('-RSB-', ']'). \
+            replace('-LCB-', '{'). \
+            replace('-RCB-', '}')
+
+    print('Loading wiki articles...')
+    articles = {}
+    for i in range(109):
+        wid = str(i + 1).zfill(3)
+        file = open(path + '/wiki-pages/wiki-' + wid + '.jsonl', 'r', encoding='utf8')
+        for line in file:
+            data = json.loads(line.rstrip('\n'))
+            lines = data['lines'].split('\n')
+            plines = []
+            for line in lines:
+                slines = line.split('\t')
+                if len(slines) > 1 and slines[0].isnumeric():
+                    plines.append(slines[1])
+                else:
+                    plines.append('')
+
+            lines = plines
+            data['id'] = unidecode(data['id'])
+            articles[data['id']] = lines
+
+    print('Preprocessing dataset...')
+    files = ['train', 'dev']
+    for file in files:
+        fo = open(path + '/facts_{}.jsonl'.format(file), 'w+', encoding='utf8')
+        with open(path + '/{}.jsonl'.format(file), encoding='utf8') as f:
+            for line in f:
+                data = json.loads(line.rstrip('\n'))
+                claim = data['claim']
+                label = data['label']
+                evidence = []
+                sents = []
+                if label != 'NOT ENOUGH INFO':
+                    for evidence_list in data['evidence']:
+                        evidence_sents = []
+                        extra_left = []
+                        extra_right = []
+                        evidence_words = 0
+                        for e in evidence_list:
+                            e[2] = unidecode(e[2])
+                            if e[2] in articles and len(articles[e[2]]) > e[3]:
+                                eid = e[3]
+                                article = articles[e[2]]
+                                evidence_sents.append(replace_symbols(article[eid].replace('  ', ' ').strip()))
+                                evidence_words += len(article[eid].split(' '))
+                                left = []
+                                for i in range(0, eid):
+                                    left.append(replace_symbols(article[i]))
+                                extra_left.append(left)
+                                right = []
+                                for i in range(eid + 1, len(article)):
+                                    right.append(replace_symbols(article[i]))
+                                extra_right.append(right)
+                            else:
+                                evidence_sents = []
+                                break
+                        evidence_text = []
+                        for i in range(len(evidence_sents)):
+                            for j in range(len(extra_left[i])):
+                                xlen = len(extra_left[i][j].split(' '))
+                                if evidence_words + xlen < 254:
+                                    evidence_text.append(extra_left[i][j])
+                                    evidence_words += xlen
+                            evidence_text.append(evidence_sents[i])
+                            for j in range(len(extra_right[i])):
+                                xlen = len(extra_right[i][j].split(' '))
+                                if evidence_words + xlen < 254:
+                                    evidence_text.append(extra_right[i][j])
+                                    evidence_words += xlen
+
+                        if len(evidence_text) > 0:
+                            evidence_text = unidecode(' '.join(evidence_text))
+                            evidence_text = ' '.join(evidence_text.split()).strip()
+                            evidence.append(evidence_text)
+                            sents.append(' '.join(evidence_sents))
+
+                if label == 'NOT ENOUGH INFO' or len(evidence) > 0:
+                    fo.write(json.dumps({
+                        'statement': unidecode(replace_symbols(claim)),
+                        'label': label,
+                        'evidence_long': list(set(evidence)),
+                        'evidence': list(set(sents))
+                    }) + '\n')
+        fo.close()
+
+
+def preprocess_squad_dataset(path, output):
+    # Read dataset
+    with open(path) as f:
+        dataset = json.load(f)
+
+    # Iterate and write question-answer pairs
+    with open(output, 'w+') as f:
+        for article in dataset['data']:
+            for paragraph in article['paragraphs']:
+                context = paragraph['context']
+                for qa in paragraph['qas']:
+                    question = qa['question']
+                    answers = [a['text'] for a in qa['answers']]
+                    min_answer = ''
+                    min_length = 1000
+                    for answer in answers:
+                        if len(answer) < min_length:
+                            min_answer = answer
+                            min_length = len(answer)
+                    if min_answer == '':
+                        continue
+                    f.write(json.dumps({'question': question, 'answer': min_answer, 'evidence': context}))
+                    f.write('\n')
+
+
+def preprocess_nq_dataset():
+    def _clean_token(token):
+        return re.sub(u" ", "_", token["token"])
+
+    def get_text_blocks(html):
+        w = WebParser()
+        html = '<div>{}</div>'.format(html)
+        w.feed(html)
+        blocks = w.get_blocks()
+        return blocks
+
+    files = ['train']
+    for file in files:
+        of = open('../data/qa/nq/{}_all_h.jsonl'.format(file), 'w+')
+        with open('../data/qa/nq/{}.jsonl'.format(file), encoding='utf8') as f:
+            for line in tqdm(f):
+                data = json.loads(line.rstrip('\n'))
+                question = data['question_text']
+
+                if file == 'dev':
+                    data['document_text'] = " ".join([_clean_token(t) for t in data["document_tokens"]])
+                doc_tokens = data['document_text'].replace('*', '').split(' ')
+
+                short_answer = ''
+                if len(data['annotations'][0]['short_answers']) > 0:
+                    short_answer_info = data['annotations'][0]['short_answers'][0]
+                    doc_tokens.insert(short_answer_info['start_token'], '*')
+                    doc_tokens.insert(short_answer_info['end_token'] + 1, '*')
+                    short_answer = ' '.join(
+                        doc_tokens[short_answer_info['start_token'] + 1:short_answer_info['end_token'] + 1])
+                    short_answer = ' '.join(get_text_blocks(short_answer))
+
+                long_answer_info = data['annotations'][0]['long_answer']
+                long_answer = ' '.join(doc_tokens[long_answer_info['start_token']:long_answer_info['end_token']])
+                long_answer = ' '.join(get_text_blocks(long_answer))
+
+                if long_answer == '':
+                    continue
+
+                example = {
+                    'question': question,
+                    'answer': short_answer,
+                    'evidence': long_answer
+                }
+
+                of.write(json.dumps(example) + '\n')
+
+
+def preprocess_msmarco():
+    files = ['train', 'dev']
+    for file in files:
+        fo = open('../data/qa/msmarco/{}_all.jsonl'.format(file), 'w+')
+        with open('../data/qa/msmarco/{}.jsonl'.format(file), 'r') as f:
+            for line in f:
+                data = json.loads(line.rstrip('\n'))
+                q = data['query']
+                a = data['answers'][0]
+                evidence = ''
+                for passage in data['passages']:
+                    if passage['is_selected'] == 1:
+                        evidence = passage['passage_text']
+                        break
+
+                if '_' in q:
+                    continue
+
+                example = {
+                    'question': q,
+                    'answer': a,
+                    'evidence': evidence
+                }
+
+                fo.write(json.dumps(example) + '\n')
+
+
+def create_qa_ranking_dataset():
+    files = ['train', 'dev']
+    for file in files:
+        relevant = 0
+        irrelevant = 0
+        max_irrelevant = 10 if file == 'train' else 1
+        fo = open('../data/qa_ranking_large/{}.jsonl'.format(file), 'w+')
+        with open('../data/qa/msmarco/{}.jsonl'.format(file), 'r') as f:
+            for line in f:
+                data = json.loads(line.rstrip('\n'))
+                added_irrelevant = 0
+                for passage in data['passages']:
+                    example = {
+                        'question': data['query'],
+                        'answer': passage['passage_text'],
+                        'relevant': passage['is_selected']
+                    }
+                    if passage['is_selected'] == 1 or added_irrelevant < max_irrelevant:
+                        fo.write(json.dumps(example) + '\n')
+                        if passage['is_selected'] == 1:
+                            relevant += 1
+                        else:
+                            added_irrelevant += 1
+                            irrelevant += 1
+
+        print(relevant, irrelevant)
+
+
+def preprocess_qa_nli_dataset(path):
+    files = ['dev', 'train']
+    for file in files:
+        of = open(path + '/' + file + '_.tsv', 'w+')
+        with open(path + '/' + file + '.tsv', 'r') as f:
+            for line in f:
+                line = line.rstrip()
+                s = line.split('\t')
+                if s[2].endswith('.'):
+                    s[2] = s[2][:len(s[2]) - 1] + '?'
+                if not s[2].endswith('?'):
+                    s[2] += ' ?'
+                input_str = s[2] + ' ' + s[3]
+                output_str = s[4]
+                of.write(input_str + '\t' + output_str + '\n')
+        of.close()
+
+
+def get_pair_input(tokenizer, sent1, sent2, max_len=256):
+    text = "[CLS] {} [SEP] {} [SEP]".format(sent1, sent2)
+
+    tokenized_text = tokenizer.tokenize(text)[:max_len]
+    indexed_tokens = tokenizer.encode(text)[:max_len]
+
+    segments_ids = []
+    sep_flag = False
+    for i in range(len(tokenized_text)):
+        if tokenized_text[i] == '[SEP]' and not sep_flag:
+            segments_ids.append(0)
+            sep_flag = True
+        elif sep_flag:
+            segments_ids.append(1)
+        else:
+            segments_ids.append(0)
+    return indexed_tokens, segments_ids
+
+
+def build_batch(tokenizer, text_list, max_len=256):
+    token_id_list = []
+    segment_list = []
+    attention_masks = []
+    longest = -1
+
+    for pair in text_list:
+        sent1, sent2 = pair
+        ids, segs = get_pair_input(tokenizer, sent1, sent2, max_len=max_len)
+        if ids is None or segs is None:
+            continue
+        token_id_list.append(ids)
+        segment_list.append(segs)
+        attention_masks.append([1] * len(ids))
+        if len(ids) > longest:
+            longest = len(ids)
+
+    if len(token_id_list) == 0:
+        return None, None, None
+
+    # padding
+    assert (len(token_id_list) == len(segment_list))
+    for ii in range(len(token_id_list)):
+        token_id_list[ii] += [0] * (longest - len(token_id_list[ii]))
+        attention_masks[ii] += [1] * (longest - len(attention_masks[ii]))
+        segment_list[ii] += [1] * (longest - len(segment_list[ii]))
+
+    return token_id_list, segment_list, attention_masks
+
+
+"""
+class QAEncoderBatchGenerator(Sequence):
+    def __init__(self, ids_to_paragraphs: List, ids_to_queries: Dict, qrels_train: List, tokenizer,
+                 max_question_length=64, max_answer_length=512, batch_size=64):
+        np.random.seed(42)
+        self.ids_to_paragraphs = ids_to_paragraphs
+        self.ids_to_queries = ids_to_queries
+        self.qrels_train = qrels_train
+        self.tokenizer = tokenizer
+        self.max_question_length = max_question_length
+        self.max_answer_length = max_answer_length
+        self.batch_size = batch_size
+        self.on_epoch_end()
+
+    def __len__(self):
+        return int(np.floor(len(self.qrels_train) / self.batch_size))
+
+    def __getitem__(self, index):
+        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]
+
+        question_texts = []
+        answer_texts = []
+
+        for index in indexes:
+            query_id = self.qrels_train[index][0]
+            answer_id = self.qrels_train[index][1]
+            question_texts.append(self.ids_to_queries[query_id])
+            answer_texts.append(self.ids_to_paragraphs[answer_id])
+
+        question_inputs = get_inputs_text(texts=question_texts, tokenizer=self.tokenizer,
+                                          max_length=self.max_question_length)
+        answer_inputs = get_inputs_text(texts=answer_texts, tokenizer=self.tokenizer,
+                                        max_length=self.max_answer_length)
+
+        x = {
+            'question_input_word_ids': question_inputs[0],
+            'question_input_masks': question_inputs[1],
+            'question_input_segments': question_inputs[2],
+            'answer_input_word_ids': answer_inputs[0],
+            'answer_input_masks': answer_inputs[1],
+            'answer_input_segments': answer_inputs[2]
+        }
+        y = np.zeros((self.batch_size, self.batch_size))
+
+        return x, y
+
+    def on_epoch_end(self, logs=None):
+        self.indexes = np.arange(len(self.qrels_train))
+        np.random.shuffle(self.indexes)
+"""
+
+
+def create_qa_encoder_pretraining_dataset(dataset_file, output_file, total_samples):
+    print('Loading paragraphs in memory...')
+
+    articles = []
+    file = open(dataset_file, 'r', encoding='utf8')
+    while True:
+        line = file.readline()
+        if not line:
+            break
+        line = line.rstrip('\n')
+        data = json.loads(line)
+        articles.append(data['paragraphs'])
+
+    print('Loaded paragraphs in memory')
+
+    o = open(output_file, 'w+')
+
+    added = 0
+    while added < total_samples:
+        index = randint(0, len(articles) - 1)
+        article = articles[index]
+        if len(article) < 2:
+            continue
+        if added % 2 == 0:  # Body-First-Selection
+            paragraph = article[0]
+            random_question_sentence_id = randint(0, len(paragraph) - 1)
+            question = paragraph[random_question_sentence_id]
+            random_answer_paragraph_id = randint(1, len(article) - 1)
+            answer = ' '.join(article[random_answer_paragraph_id])
+        else:  # Inverse-Cloze-Task
+            random_paragraph_id = randint(0, len(article) - 1)
+            paragraph = article[random_paragraph_id]
+            random_question_sentence_id = randint(0, len(paragraph) - 1)
+            question = paragraph[random_question_sentence_id]
+            choice = random()
+            if choice < 0.9:
+                answer = ' '.join([paragraph[j] for j in range(len(paragraph)) if j != random_question_sentence_id])
+            else:
+                answer = ' '.join(paragraph)
+
+        o.write(json.dumps({'s1': question, 's2': answer}) + '\n')
+
+        added += 1
+        if added % 10000 == 0:
+            print(added)
+
+    o.close()
+
+
+def create_random_wiki_paragraphs(dataset_file, output_file, total_samples, max_length=510):
+    print('Loading paragraphs in memory...')
+
+    articles = []
+    file = open(dataset_file, 'r', encoding='utf8')
+    while True:
+        line = file.readline()
+        if not line:
+            break
+        line = line.rstrip('\n')
+        data = json.loads(line)
+        articles.append(data['paragraphs'])
+
+    print('Loaded paragraphs in memory')
+
+    o = open(output_file, 'w+')
+    added = 0
+    while added < total_samples:
+        index = randint(0, len(articles) - 1)
+        article = articles[index]
+        if len(article) < 1:
+            continue
+        random_paragraph_id = randint(0, len(article) - 1)
+        paragraph = article[random_paragraph_id]
+
+        p = []
+        words = 0
+        i = 0
+        while i < len(paragraph):
+            nwords = len(paragraph[i].split(' '))
+            if nwords + words >= max_length:
+                break
+            p.append(paragraph[i])
+            i += 1
+
+        o.write(json.dumps({'text': ' '.join(p)}) + '\n')
+        added += 1
+        if added % 10000 == 0:
+            print(added)
+    o.close()
+
+
+def load_unsupervised_dataset(dataset_file):
+    print('Loading dataset...')
+    x = pickle.load(open(dataset_file, "rb"))
+    print('Done')
+    return x, len(x[0])
+
+
+def create_entailment_dataset(dataset_file, tokenizer, output_file):
+    x = [[] for _ in range(6)]
+    y = []
+
+    with open(dataset_file, encoding='utf8') as file:
+        for j, line in enumerate(file):
+            if j % 10000 == 0:
+                print(j)
+            line = line.rstrip('\n')
+            sample = json.loads(line)
+            label = sample['label']
+            s1 = tokenize(text=sample['s1'],
+                          max_length=256,
+                          tokenizer=tokenizer)
+            s2 = tokenize(text=sample['s2'],
+                          max_length=256,
+                          tokenizer=tokenizer)
+            example = s1 + s2
+            for i in range(6):
+                x[i].append(example[i])
+            if label == 'entailment':
+                label = [1.0, 0.0, 0.0]
+            elif label == 'contradiction':
+                label = [0.0, 1.0, 0.0]
+            else:
+                label = [0.0, 0.0, 1.0]
+            y.append(np.asarray(label, dtype='float32'))
+
+    for i in range(6):
+        x[i] = np.asarray(x[i])
+
+    y = np.asarray(y)
+    data = [x, y]
+
+    pickle.dump(data, open(output_file, "wb"), protocol=4)
+
+
+def load_supervised_dataset(dataset_file):
+    print('Loading dataset...')
+    d = pickle.load(open(dataset_file, "rb"))
+    print('Done')
+    return d[0], d[1]
+
+
+def preprocess_sg():
+    fo = open('../data/qa/squad/dev_sg_sq.txt', 'w+')
+    f = open('../data/qa/squad/dev.jsonl', 'r')
+    for line in f:
+        d = json.loads(line.rstrip('\n'))
+        q = d['question']
+        a = d['answer']
+
+        if q.endswith('.'):
+            q = q[:len(q) - 1] + '?'
+        if not q.endswith('?'):
+            q += ' ?'
+
+        fo.write(q + ' ' + a + '\n')
+
+
+def preprocess_msmarco_():
+    fo = open('../data/qa/msmarco/dev_sg.tsv', 'w+')
+    with open('../data/qa/msmarco/dev.jsonl', 'r') as f:
+        for line in f:
+            data = json.loads(line.rstrip('\n'))
+            if data['wellFormedAnswers'] == '[]' or data['query_type'] == 'DESCRIPTION':
+                continue
+            q = data['query']
+            a = data['answers'][0]
+            fa = data['wellFormedAnswers'][0]
+
+            if len(a) > len(q):
+                continue
+
+            if q.endswith('.'):
+                q = q[:len(q) - 1] + '?'
+            if not q.endswith('?'):
+                q += ' ?'
+
+            if '_' in q:
+                continue
+
+            fo.write('{} {}\t{}\n'.format(q, a, fa))
+
+
+def create_qnli_nq():
+    files = ['train', 'dev']
+    for file in files:
+        o = open('../data/qnli/{}_nq.jsonl'.format(file), 'w+')
+
+        fs = open('../data/qnli/{}.txt'.format(file))
+        statements = [s.rstrip('\n') for s in fs]
+
+        fd = open('../data/qa/nq/{}_nli.jsonl'.format(file))
+        for i, line in enumerate(fd):
+            d = json.loads(line.rstrip('\n'))
+            o.write(json.dumps({
+                'statement': statements[i],
+                'question': d['question'],
+                'answer': d['answer'],
+                'evidence': d['evidence']
+            }) + '\n')
+
+
+def create_qnli_ms():
+    files = ['train', 'dev']
+    for file in files:
+        o = open('../data/qnli/{}_ms.jsonl'.format(file), 'w+')
+        with open('../data/qa/msmarco/{}.jsonl'.format(file), 'r') as f:
+            for line in f:
+                data = json.loads(line.rstrip('\n'))
+                if data['wellFormedAnswers'] == '[]' or data['query_type'] == 'DESCRIPTION':
+                    continue
+                q = data['query']
+                a = data['answers'][0]
+                fa = data['wellFormedAnswers'][0]
+
+                evidence = ''
+                for p in data['passages']:
+                    if p['is_selected'] == 1:
+                        evidence = p['passage_text']
+                if evidence == '' or '_' in q:
+                    continue
+
+                o.write(json.dumps({
+                    'statement': fa,
+                    'question': q,
+                    'answer': a,
+                    'evidence': evidence
+                }) + '\n')
+
+
+def create_qnli_ms_2():
+    files = ['train', 'dev']
+    for file in files:
+        o = open('../data/qnli/{}_ms_nw.jsonl'.format(file), 'w+')
+        fs = open('../data/qa/msmarco/statements_{}_ms.txt'.format(file))
+        statements = [s.rstrip('\n') for s in fs]
+        with open('../data/qa/msmarco/{}.jsonl'.format(file), 'r') as f:
+            n = 0
+            for line in f:
+                data = json.loads(line.rstrip('\n'))
+                if data['wellFormedAnswers'] != '[]' or data['query_type'] == 'DESCRIPTION':
+                    continue
+                q = data['query']
+                a = data['answers'][0]
+
+                if a == 'No Answer Present.' or '_' in q:
+                    continue
+
+                evidence = ''
+                for p in data['passages']:
+                    if p['is_selected'] == 1:
+                        evidence = p['passage_text']
+
+                if '\n' in q or '\n' in a:
+                    nl = (q + a).count('\n') + 1
+                    print(nl)
+                    n += nl
+                    continue
+
+                o.write(json.dumps({
+                    'statement': statements[n],
+                    'question': q,
+                    'answer': a,
+                    'evidence': evidence
+                }) + '\n')
+
+                n += 1
+
+
+def create_msmarco_sg():
+    files = ['train', 'dev']
+    for file in files:
+        o = open('../data/qa/msmarco/{}_ms_sg.txt'.format(file), 'w+')
+        with open('../data/qa/msmarco/{}.jsonl'.format(file), 'r') as f:
+            for line in f:
+                data = json.loads(line.rstrip('\n'))
+                if data['wellFormedAnswers'] != '[]' or data['query_type'] == 'DESCRIPTION':
+                    continue
+                q = data['query']
+                a = data['answers'][0]
+
+                if a == 'No Answer Present.' or '_' in q:
+                    continue
+
+                if q.endswith('.'):
+                    q = q[:len(q) - 1] + '?'
+                if not q.endswith('?'):
+                    q += ' ?'
+
+                o.write(q + ' ' + a + '\n')
+                """
+                o.write(json.dumps({
+                    'statement': fa,
+                    'question': q,
+                    'answer': a,
+                    'evidence': evidence
+                }) + '\n')
+                """
+
+
+def create_qar():
+    files = ['train', 'dev']
+    for file in files:
+        o = open('../data/qar/{}2.jsonl'.format(file), 'w+')
+
+        f = open('../data/qa/nq/{}_all.jsonl'.format(file))
+        for line in f:
+            d = json.loads(line.rstrip('\n'))
+            if d['evidence'] == '':
+                continue
+            o.write(json.dumps({
+                'question': d['question'],
+                'answer': d['evidence'],
+                'source': 'natural-questions'
+            }) + '\n')
+
+        f = open('../data/qa/msmarco/{}_all.jsonl'.format(file))
+        for line in f:
+            d = json.loads(line.rstrip('\n'))
+            if d['evidence'] == '':
+                continue
+            o.write(json.dumps({
+                'question': d['question'],
+                'answer': d['evidence'],
+                'source': 'msmarco'
+            }) + '\n')
+
+
+"""
+question_inputs = get_inputs_text(texts=question_texts, tokenizer=tokenizer,
+                                      max_length=max_question_length)
+    answer_inputs = get_inputs_text(texts=answer_texts, tokenizer=tokenizer,
+                                    max_length=max_answer_length)
+
+    x = {
+        'question_input_word_ids': question_inputs[0],
+        'question_input_masks': question_inputs[1],
+        'question_input_segments': question_inputs[2],
+        'answer_input_word_ids': answer_inputs[0],
+        'answer_input_masks': answer_inputs[1],
+        'answer_input_segments': answer_inputs[2]
+    }
+    y = np.zeros((self.batch_size, self.batch_size))
+"""
+
+"""
+create_random_wiki_paragraphs('../data/wikipedia/wiki_paragraphs.jsonl',
+                              '../data/random_wiki_paragraphs.jsonl',
+                              total_samples=1000000,
+                              max_length=510)
+"""
+
+# preprocess_nq_dataset()

models/dense_retriever.py

@@ -0,0 +1,47 @@
+import h5py
+import logging
+import pickle
+import sqlite3
+import struct
+
+import numpy as np
+
+from models.vector_index import VectorIndex
+
+import random
+
+
+class DenseRetriever:
+    def __init__(self, model, db_path, batch_size=64, use_gpu=False, debug=False):
+        self.model = model
+        self.vector_index = VectorIndex(768)
+        self.batch_size = batch_size
+        self.use_gpu = use_gpu
+        self.db = sqlite3.connect(db_path)
+        self.db.row_factory = sqlite3.Row
+        self.debug = debug
+
+    def load_pretrained_index(self, path):
+        self.vector_index.load(path)
+
+    def populate_index(self, table_name):
+        cur = self.db.cursor()
+        query = f'SELECT * FROM {table_name} ORDER BY idx' if not self.debug \
+            else f'SELECT * FROM {table_name} ORDER BY idx LIMIT 1000'
+        for r in cur.execute(query):
+            e = r['encoded']
+            v = [np.float32(struct.unpack('f', e[i*4:(i+1)*4])[0]) for i in range(int(len(e)/4))]
+            self.vector_index.index.add(np.ascontiguousarray([v]))
+            print(f"\rAdded {self.vector_index.index.ntotal} vectors", end='')
+        print()
+        logging.info("Finished adding vectors.")
+
+    def search(self, queries, limit=1000, probes=512, min_similarity=0):
+        query_vectors = self.model.encode(queries, batch_size=self.batch_size)
+        ids, similarities = self.vector_index.search(query_vectors, k=limit, probes=probes)
+        results = []
+        for j in range(len(ids)):
+            results.append([
+                (ids[j][i], similarities[j][i]) for i in range(len(ids[j])) if similarities[j][i] > min_similarity
+            ])
+        return results

models/encoder/0_BERT/config.json

@@ -0,0 +1,20 @@
+{
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "output_past": true,
+  "pad_token_id": 0,
+  "type_vocab_size": 2,
+  "vocab_size": 28996
+}

models/encoder/0_BERT/sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 256,
+  "do_lower_case": false
+}

models/encoder/0_BERT/special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

models/encoder/0_BERT/tokenizer_config.json

@@ -0,0 +1 @@
+{"do_lower_case": false, "max_len": 512, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

models/encoder/1_Pooling/config.json

@@ -0,0 +1,7 @@
+{
+  "word_embedding_dimension": 768,
+  "pooling_mode_cls_token": false,
+  "pooling_mode_mean_tokens": true,
+  "pooling_mode_max_tokens": false,
+  "pooling_mode_mean_sqrt_len_tokens": false
+}

models/encoder/config.json

@@ -0,0 +1,3 @@
+{
+  "__version__": "1.0"
+}

models/encoder/modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "0_BERT",
+    "type": "sentence_transformers.models.BERT"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  }
+]

models/encoder/stats.csv

@@ -0,0 +1,220 @@
+0.5779840120767146,0.8015233128623872,0.8583559199917659,0.8996466188630048,0.9569938587161629,0.6784459557674578
+0.5916046248327443,0.817116684392905,0.8708958040278588,0.9097505746732082,0.9613682368682883,0.6921505611721683
+0.6190860122825677,0.8418190551343191,0.8930078567262497,0.927230932857584,0.970237074141421,0.7180167815120542
+0.6270628195011494,0.8501046419871685,0.8993378392287371,0.9313308402236937,0.9730847085463341,0.7257908699942954
+0.6326380073420935,0.8536212989329948,0.9017051497581227,0.9325831131848904,0.9728273921844444,0.7306655399664306
+0.6432565958760764,0.8622328198442378,0.9086698459532714,0.9393248018664013,0.9762754314337667,0.740502524041259
+0.6463272377946272,0.8637252547431983,0.9105739870312554,0.94061138367585,0.9756750265893573,0.7430413705788191
+0.6501698287988472,0.8689573541016228,0.9141592616735856,0.9429100765087316,0.9773390057295777,0.7468761817476934
+0.6538580299859333,0.8707757230589769,0.9151542182728926,0.9432531649912512,0.9766699831886644,0.7501280666772407
+0.6536007136240436,0.8742066078841734,0.9184993309774591,0.9457920197618966,0.9773218513054517,0.751086942266862
+0.6515593371530518,0.8723539300785672,0.9152399903935224,0.9439908052286685,0.9773561601537036,0.7486908765241613
+0.6565512745737125,0.871650598689402,0.9158060863896799,0.9442824304388102,0.9770645349435619,0.7520715039306786
+0.6568600542079802,0.8752358733317323,0.919545750849144,0.9464781967269359,0.9773733145778296,0.7533833000405535
+0.6647682437300579,0.8794387072425979,0.922376230829931,0.9491371324664631,0.9790716025663019,0.7595162032832385
+0.6617490650838851,0.8790270010635743,0.919528596425018,0.9472158369643531,0.9780938003911208,0.7577154555045574
+0.6612172779359797,0.8765739184135588,0.9201290012694274,0.945397468006999,0.9761381960407589,0.756179860041224
+0.664836861426562,0.8797989501492435,0.9218787525302775,0.9485367276220538,0.9781452636634989,0.7597794971475365
+0.6645795450646722,0.8795073249391018,0.9216557450166398,0.9480220948982743,0.9778707928774831,0.7592882243172064
+0.6696229457577109,0.8815658558342196,0.9232682608844821,0.9485367276220538,0.9783511167530106,0.7631352909333269
+0.6653686485744673,0.8806052080831647,0.9231824887638522,0.9483480289566679,0.9792259923834357,0.7605593087121496
+0.6703262771468762,0.8829210553401722,0.924383298452671,0.9496689196143685,0.9781967269358768,0.763841387028115
+0.6659004357223728,0.8810169142621882,0.9229594812502144,0.9479020139293924,0.9770816893676879,0.7606746681736507
+0.6761587813497101,0.8852025937489278,0.9259100421998834,0.9505437952447936,0.9787456685079082,0.7682947389867827
+0.6741688681510961,0.8837273132740934,0.9250866298418362,0.9504580231241637,0.9782996534806326,0.7671206271858944
+0.6687137612790338,0.882183415102755,0.9244862249974268,0.9500120080968882,0.9785569698425224,0.7634663607979358
+0.6759014649878203,0.8859745428345971,0.9264418293477888,0.9514529797234706,0.9791745291110577,0.7687004551272422
+0.6738600885168284,0.8873640511888016,0.9274367859470958,0.9514015164510927,0.9790887569904279,0.7675660379918043
+0.6721446461042303,0.8834185336398257,0.9238858201530175,0.9491885957388411,0.9771503070641918,0.7653706174723683
+0.673396919065427,0.8885991697258723,0.9287233677565444,0.9533399663773288,0.9803581843757505,0.7681142393185553
+0.677205201221395,0.8848766596905342,0.9255669537173637,0.9494802209489828,0.9775963220914674,0.7691179334267018
+0.6725906611315058,0.8832984526709439,0.9240916732425293,0.9494116032524789,0.9772189247606958,0.765669866288031
+0.679126496723505,0.8885648608776203,0.9288091398771743,0.9522249288091399,0.978900058325042,0.7717041464381332
+0.6765018698322297,0.8865749476790065,0.9267849178303085,0.9513500531787148,0.9794146910488215,0.7692736782586759
+0.6792122688441349,0.8870724259786599,0.9285689779394106,0.9536144371633444,0.9801351768621127,0.771438903068368
+0.6806360860465914,0.8879816104573369,0.9286204412117885,0.9532370398325728,0.9795690808659553,0.772461625855239
+0.677170892373143,0.8858716162898411,0.926287439530655,0.9515044429958486,0.9794661543211994,0.7698102128989501
+0.685816722132638,0.8909493258311318,0.9305588911380245,0.9536830548598484,0.9806498095858922,0.776847345661558
+0.6874292380004803,0.8926304593954781,0.9318969362198511,0.9543177685525097,0.9809757436442859,0.778282769703823
+0.6828490067588431,0.8883418533639825,0.9285003602429066,0.9527567159570454,0.9795176175935774,0.774156750244517
+0.6826946169417093,0.890726318317494,0.9309705973170481,0.954986791093423,0.980889971523656,0.7746702901532684
+0.679143651147631,0.8890451847531479,0.9278313377019933,0.9516759872371084,0.9774419322743335,0.771751324626232
+0.6816310426458984,0.8901602223213366,0.9304731190173946,0.9536144371633444,0.9792088379593097,0.7738497806623483
+0.682986242151851,0.8900229869283288,0.929066456239064,0.95205338456788,0.9785741242666484,0.7741901341300708
+0.6789549524822451,0.8901945311695887,0.9288949119978043,0.9527052526846673,0.9786084331149003,0.7717603788208348
+0.687738017634748,0.8921844443682025,0.9321027893093629,0.9541290698871239,0.9797749339554671,0.778173339236671
+0.6816653514941503,0.8897313617181871,0.9295124712663396,0.9536487460115964,0.9787456685079082,0.7736270773705498
+0.6899852471952517,0.8917212749168011,0.930095721686623,0.9531512677119429,0.9784197344495146,0.7792509607389516
+0.68392973547878,0.8919271280063128,0.9309191340446701,0.9545922393385254,0.9804096476481284,0.7758120806040231
+0.687772326483,0.8929735478779978,0.9319483994922291,0.9544893127937695,0.980907125947782,0.7784736944841473
+0.6842042062647957,0.8922359076405805,0.9316224654338354,0.9548838645486671,0.980804199403026,0.7761278134585238
+0.6846159124438193,0.8934538717535252,0.9321199437334888,0.9548495557004152,0.9802381034068687,0.7768488500009861
+0.6882354959344015,0.8937626513877929,0.9338010772978351,0.9569423954437849,0.9811815967337977,0.7791075653748736
+0.6855250969224963,0.8929049301814939,0.9323086423988747,0.9555871959378324,0.9802895666792466,0.77714622677133
+0.6878752530277559,0.8921501355199506,0.9322057158541188,0.955021099941675,0.9800494047414828,0.778441356120615
+0.684873228805709,0.8899543692318249,0.9298212509006073,0.9535458194668405,0.9789686760215459,0.7759056803242941
+0.684890383229835,0.8920815178234467,0.931982708340481,0.9547980924280371,0.9799979414691049,0.7768483470243297
+0.6939479191683535,0.8962157340378083,0.9336466874807012,0.9556901224825882,0.9813531409750574,0.7834821941604841
+0.6864514358252993,0.8933680996328953,0.9311936048306858,0.9543863862490136,0.9796376985624593,0.7777730592023087
+0.6901567914365114,0.8946889902905959,0.9330634370604178,0.9547980924280371,0.9795347720177033,0.7806016928782297
+0.6914776820942121,0.8939513500531787,0.9313479946478197,0.9534600473462106,0.9786598963872782,0.7812826686600813
+0.6920266236662436,0.894637527018218,0.9334408343911895,0.9557758946032182,0.9812502144303016,0.781829703594363
+0.6930215802655505,0.8951521597419975,0.9326517308813943,0.9553298795759426,0.9809242803719079,0.7825955960750406
+0.6919923148179915,0.8950149243489895,0.9326517308813943,0.9554499605448246,0.9804096476481284,0.781802873512625
+0.6883727313274094,0.8899543692318249,0.9277970288537414,0.9525508628675335,0.9770816893676879,0.7779526954120721
+0.6861769650392836,0.8913781864342813,0.9306275088345284,0.9539575256458641,0.9791059114145538,0.7770877463274903
+0.6918550794249837,0.8962328884619344,0.9347617250488901,0.9580059697395958,0.9829313479946479,0.7821665815849399
+0.6930901979620544,0.8964044327031941,0.9333550622705595,0.9554671149689505,0.9809414347960339,0.7827637999409208
+0.6845816035955673,0.8918928191580608,0.9293237726009538,0.9518303770542423,0.9778364840292312,0.775727886193286
+0.6941023089854873,0.8980684118434145,0.9342470923251106,0.9558788211479741,0.981833464850585,0.7839931362022297
+0.689573541016228,0.8941572031426905,0.9325831131848904,0.9552269530311868,0.9811815967337977,0.7801137155274634
+0.6899680927711257,0.891807047037431,0.9305760455621505,0.9524136274745256,0.9781452636634989,0.7794242917082135
+0.6926441829347789,0.895598174769273,0.9327889662744022,0.9556215047860843,0.9811472878855457,0.7823044597371529
+0.6946855594057707,0.8966789034892099,0.9346244896558823,0.9555871959378324,0.9814217586715615,0.7841933350778004
+0.690654269736165,0.8944659827769582,0.9323257968230007,0.9552441074553127,0.9805297286170104,0.7806231726203139
+0.6921810134833773,0.895649638041651,0.9332006724534258,0.955913129996226,0.9812159055820496,0.7820928716391452
+0.6918207705767317,0.895529557072769,0.9342985555974886,0.9558445122997221,0.9821250900607267,0.7823101370916966
+0.6965382372113769,0.899440765773493,0.9366658661268741,0.9581775139808557,0.9829313479946479,0.7861867148449412
+0.6914433732459602,0.8950320787731156,0.9328232751226542,0.9559817476927299,0.9814389130956874,0.7817211414999055
+0.6889731361718188,0.8940542765979346,0.931999862764607,0.9545922393385254,0.9805983463135143,0.7797127544519308
+0.6913061378529523,0.8949634610766116,0.9325144954883865,0.9561189830857378,0.9811815967337977,0.781525423126717
+0.6936048306858339,0.896473050399698,0.9326002676090164,0.9560846742374859,0.9814903763680654,0.7835110543483639
+0.6941023089854873,0.8942258208391944,0.9313308402236937,0.9537859814046042,0.978917212749168,0.7828756585433598
+0.6987168490753766,0.8988403609290836,0.9345901808076302,0.9566336158095172,0.9807527361306481,0.787165665870235
+0.69545750849144,0.8965416680962021,0.9326860397296463,0.9552269530311868,0.9796376985624593,0.7842072813024578
+0.6983737605928569,0.8981370295399184,0.9356537551034412,0.9574055648951865,0.9815075307921913,0.7869324771044554
+0.696383847394243,0.8968676021545957,0.9340583936597249,0.9562562184787456,0.9800150958932309,0.7851641911299555
+0.6974988849624318,0.8990119051703435,0.9364771674614883,0.9584348303427453,0.9822794798778605,0.7869184646060629
+0.6960750677599753,0.8977939410573987,0.9344872542628744,0.9556729680584622,0.979912169348475,0.7852065186352417
+0.6984938415617388,0.8997838542560126,0.936785947095756,0.9581775139808557,0.9815933029128212,0.787530466944267
+0.6984766871376128,0.8989261330497135,0.935876762617079,0.9571997118056746,0.9802037945586166,0.7871252651258276
+0.6960922221841013,0.8987888976567057,0.9356022918310632,0.956959549867911,0.9812330600061756,0.7855466480210304
+0.6972415686005421,0.8972278450612413,0.9343157100216146,0.9560675198133599,0.9810786701890417,0.7856690263774906
+0.6998661954918174,0.8990805228668474,0.9345901808076302,0.9566679246577693,0.9809242803719079,0.7876728318822217
+0.6978248190208255,0.8990290595944694,0.9359968435859608,0.9581260507084777,0.9816276117610732,0.7868541067616021
+0.6961265310323532,0.8981198751157924,0.933852540570213,0.9557415857549663,0.9791402202628058,0.7853667716098274
+0.7032113081963839,0.9011047449137133,0.9365972484303702,0.9581088962843517,0.9811472878855457,0.7905386234241963
+0.7025422856554705,0.8997323909836347,0.9361169245548427,0.9574227193193124,0.9810958246131677,0.7898037355362394
+0.6948399492229046,0.8977253233608947,0.9346587985041342,0.9565993069612653,0.9812845232785535,0.7844807566485703
+0.6989227021648883,0.8963358150066902,0.9337839228737091,0.9558959755721,0.9803067211033726,0.7864312940359843
+0.6998661954918174,0.8990462140185954,0.934796033897142,0.9556215047860843,0.979912169348475,0.7876511656582856
+0.6971729509040382,0.8959412632517927,0.9324115689436305,0.954008988918242,0.9792088379593097,0.7851343541574478
+0.6985967681064946,0.9002984869797921,0.9374892784849213,0.9583319037979895,0.9815933029128212,0.7875866297739705
+0.7022506604453288,0.8989947507462175,0.9346073352317563,0.9563419905993755,0.9803410299516245,0.7894850593461494
+0.699231481799156,0.8983600370535562,0.9344700998387484,0.9559817476927299,0.979946478196727,0.7872299076283827
+0.7014787113596597,0.9008817374000755,0.936820255944008,0.9582632861014856,0.981713383881703,0.7897078452703377
+0.6975846570830617,0.8972621539094933,0.933989775963221,0.9556558136343363,0.9795519264418293,0.7858221981230687
+0.6996260335540536,0.8993721480769891,0.9361855422513466,0.9576457268329502,0.980855662675404,0.7881251554498462
+0.6976361203554396,0.9005729577658078,0.936785947095756,0.9582804405256116,0.9816276117610732,0.7872117895320041
+0.7015816379044155,0.8990119051703435,0.9352077400761657,0.956856623323155,0.9794146910488215,0.7890413758536403
+0.7018217998421793,0.9008302741276976,0.9366315572786221,0.9581432051326036,0.9818677736988369,0.7899203650821491
+0.705956016056541,0.9030775036882012,0.9381583010258345,0.9595327134868082,0.981816310426459,0.7932645670386803
+0.7033828524376436,0.9035063642913508,0.9389817133838817,0.9598414931210759,0.9826568772086322,0.7920589796937257
+0.7020276529316911,0.9032147390812091,0.9387415514461179,0.9598586475452019,0.9823137887261124,0.7907239281404607
+0.7005523724568566,0.9007445020070676,0.9359968435859608,0.9564620715682575,0.9800665591656088,0.7889600463763118
+0.7032284626205099,0.9015336055168628,0.9372491165471575,0.957834425498336,0.9813188321268055,0.7909847811260302
+0.7028853741379902,0.9019624661200123,0.9383813085394723,0.9585549113116273,0.9822108621813566,0.7910605299160958
+0.7070195903523518,0.9029574227193193,0.9383641541153463,0.957851579922462,0.9813359865509315,0.7936546808268249
+0.7017531821456754,0.901276289154973,0.936820255944008,0.9575771091364463,0.9807870449789,0.7899696413107273
+0.7035029334065256,0.9027687240539335,0.9378323669674409,0.9588979997941469,0.9815761484886952,0.7913691735211271
+0.70689950938347,0.9027858784780595,0.9388101691426218,0.9591381617319107,0.981833464850585,0.7936974704832799
+0.703039763955124,0.903300511201839,0.9389817133838817,0.9593783236696745,0.9823137887261124,0.7913347277016449
+0.7016502556009194,0.9015336055168628,0.9378666758156928,0.9586235290081312,0.9816790750334511,0.7902131444981638
+0.7057158541187772,0.9025971798126737,0.9373177342436615,0.9591553161560367,0.9822794798778605,0.7929256361254804
+0.7009640786358802,0.9010875904895873,0.9357223727999451,0.9566336158095172,0.9807184272823961,0.7890500472662872
+0.7051669125467458,0.9037808350773665,0.9386043160531101,0.9593268603972964,0.9816619206093251,0.7929897654311266
+0.7071225168971078,0.9051017257350671,0.9408343911894878,0.9611280749305245,0.983703297080317,0.794594734873485
+0.708769341613202,0.9037636806532404,0.9391361032010156,0.9600130373623358,0.9819535458194668,0.794976115563586
+0.7083404810100525,0.9041582324081381,0.9384670806601022,0.9596184856074381,0.9822108621813566,0.7949226710372033
+0.7017703365698014,0.9011562081860912,0.9369574913370158,0.9579545064672179,0.9816790750334511,0.7897549305809078
+0.7036573232236594,0.9018766939993824,0.9378495213915669,0.9589323086423989,0.981850619274711,0.7916034369572444
+0.7054585377568875,0.9039523793186263,0.9384156173877243,0.9596184856074381,0.9826568772086322,0.7932860445089596
+0.706916663807596,0.903163275808831,0.938707242597866,0.9589837719147768,0.9820050090918447,0.7939775358361482
+0.7069852815040999,0.9050845713109411,0.9400624421038186,0.9609908395375167,0.9829485024187739,0.7945385075253434
+0.7062647956908087,0.903163275808831,0.9378152125433149,0.9587950732493911,0.9812673688544276,0.7934864404440743
+0.7067108107180842,0.9052904244004529,0.9403883761622123,0.9617456341990599,0.983703297080317,0.7943919250071902
+0.7090781212474697,0.9049301814938072,0.9391017943527635,0.9599958829382098,0.9820907812124747,0.7957028931092571
+0.7073626788348715,0.9060452190619961,0.9407657734929838,0.9618828695920678,0.9837890692009469,0.7950442483313808
+0.7094212097299893,0.9053075788245789,0.9405084571310941,0.9614025457165403,0.9832915909012935,0.7962251490469788
+0.7094898274264932,0.9045699385871616,0.9400109788314406,0.960699214327375,0.98279411260164,0.7961054241729613
+0.7089065770062098,0.9050159536144372,0.9404055305863382,0.9606305966308711,0.9826911860568841,0.7957211011304709
+0.7072940611383676,0.9039180704703743,0.9389302501115038,0.9589494630665248,0.9811129790372937,0.7943070084981082
+0.7095584451229973,0.905153189007445,0.9410745531272515,0.9616255532301781,0.9830857378117817,0.7966290988294025
+0.7099358424537688,0.9054448142175867,0.9414176416097711,0.9606820599032491,0.9828970391463958,0.7963943395382952
+0.7119772189247607,0.907040175661303,0.9419494287576766,0.9619686417126977,0.9832572820530415,0.7984175961150379
+0.7118056746835009,0.9052561155522009,0.9400109788314406,0.9593783236696745,0.9812845232785535,0.797624930996624
+0.7138127423062408,0.9064226163927677,0.9405084571310941,0.9602875081483515,0.9821593989089786,0.798959296397886
+0.7152194050845713,0.9079322057158541,0.9420695097265585,0.9617284797749339,0.9833258997495454,0.8005330807512997
+0.7125433149209182,0.9075548083850825,0.9420008920300545,0.9619686417126977,0.9836518338079391,0.7988329781501703
+0.7090438123992178,0.9046728651319175,0.9386043160531101,0.9594469413661784,0.9806326551617662,0.7953849481751337
+0.7113768140803514,0.9046042474354136,0.9385700072048582,0.9591038528836587,0.9800665591656088,0.7968066498706435
+0.7120801454695166,0.9053590420969568,0.9400795965279446,0.9601331183312176,0.9816962294575771,0.7977260981425256
+0.7100559234226507,0.9055648951864685,0.9402682951933303,0.9606477510549971,0.981833464850585,0.7962313583525185
+0.7126977047380519,0.9069887123889251,0.9412117885202593,0.9620201049850756,0.9824681785432463,0.7985018939595462
+0.7109307990530758,0.9065770062099016,0.940645692524102,0.9599958829382098,0.9814903763680654,0.7971671390474596
+0.7122516897107765,0.9061138367585,0.94069715579648,0.9602017360277215,0.9814217586715615,0.7978698135709755
+0.7104676296016743,0.9054962774899646,0.940645692524102,0.9609565306892648,0.9817476927299551,0.7965702130914314
+0.7113253508079733,0.9059251380931143,0.9408687000377397,0.9609222218410128,0.9826054139362541,0.7971691791256932
+0.7113253508079733,0.9059594469413662,0.9399938244073146,0.9603904346931074,0.9814903763680654,0.7971970327904135
+0.7127663224345558,0.9066799327546574,0.9418293477887947,0.9615054722612962,0.9822108621813566,0.7987212070933885
+0.7143616838782723,0.9088928534669091,0.9430816207499915,0.962654818677737,0.9831543555082856,0.8000986555854859
+0.7152365595086972,0.9098878100662161,0.9432874738395032,0.962723436374241,0.9833430541736714,0.800961004414098
+0.7179641129447284,0.9099049644903421,0.9438021065632827,0.9623460390434693,0.9831886643565375,0.802537516552073
+0.7143102206058942,0.9087384636497753,0.9414347960338971,0.9615740899578001,0.9825882595121281,0.7999251770665431
+0.7152708683569493,0.9076577349298384,0.9418121933646687,0.9613682368682883,0.9823995608467424,0.8003246145804358
+0.7156997289600988,0.9076748893539643,0.9420523553024325,0.9611452293546505,0.9821765533331046,0.8005897851903018
+0.7151507873880674,0.9085497649843894,0.9421724362713144,0.9616255532301781,0.982776958177514,0.800298538385517
+0.7177582598552167,0.9090815521322949,0.9427385322674717,0.9619000240161938,0.9830514289635297,0.8021783842434324
+0.7155453391429649,0.9086183826808935,0.9423954437849521,0.9620372594092016,0.9829485024187739,0.8008689468904545
+0.7144817648471541,0.9089614711634131,0.9423268260884482,0.961642707654304,0.9829828112670258,0.800437947894952
+0.7129893299481936,0.9073146464473187,0.9409544721583697,0.9608192952962569,0.9825539506638762,0.7985213141537784
+0.7184444368202559,0.9101965897004838,0.9442652760146842,0.9625175832847291,0.9839777678663327,0.8031311330802511
+0.7183586646996261,0.9102823618211137,0.9441280406216763,0.9624661200123512,0.9839777678663327,0.8031333976315808
+0.71842728239613,0.9080351322606101,0.9417778845164168,0.9610594572340206,0.9825882595121281,0.8024804003813029
+0.7160085085943665,0.9082924486224997,0.9421895906954404,0.9607849864480049,0.9826397227845062,0.8007421064452407
+0.7192163859059252,0.9107283768483893,0.9447799087384636,0.9631008337050125,0.9839263045939548,0.8037904528738081
+0.7172950904038151,0.9081209043812399,0.942584142450338,0.9615397811095482,0.9822966343019864,0.801610435193648
+0.7180327306412324,0.9065255429375236,0.9410745531272515,0.9600130373623358,0.9811129790372937,0.8011702379380538
+0.7188561429992795,0.9108313033931451,0.9441966583181802,0.9632209146738944,0.9834459807184273,0.8034972168782052
+0.7175009434933269,0.9095790304319484,0.9428071499639757,0.9612653103235325,0.9821079356366007,0.8020695604970132
+0.7175009434933269,0.9095104127354444,0.9429272309328576,0.9617627886231859,0.9829313479946479,0.8021815489680041
+0.7151336329639414,0.9078121247469723,0.9413490239132672,0.9610423028098947,0.9826225683603801,0.800175479488273
+0.7194565478436888,0.9100593543074759,0.942566988026212,0.9621230315298316,0.9825196418156242,0.8033165975051533
+0.7187189076062717,0.9100593543074759,0.9426527601468419,0.9621744948022095,0.9823995608467424,0.8029329362048628
+0.7217552406765705,0.9117747967200741,0.944436820255944,0.9629807527361306,0.9831715099324115,0.8051882216053523
+0.7207431296531376,0.9112087007239167,0.9439908052286685,0.962654818677737,0.982742649329262,0.8044874117887816
+0.7204171955947439,0.9098706556420901,0.9441451950458023,0.9623288846193433,0.9828970391463958,0.8041066613024709
+0.7210004460150273,0.9107969945448932,0.9447799087384636,0.962689127525989,0.9828798847222698,0.8045389583450916
+0.7219096304937044,0.9123923559886095,0.9450886883727313,0.963632620852918,0.9833258997495454,0.8057297371947867
+0.7207431296531376,0.9095447215836965,0.9431845472947473,0.9618828695920678,0.9823309431502385,0.8041280599542056
+0.7218753216454523,0.9117919511442001,0.9450372251003534,0.9630665248567606,0.9829656568428998,0.8054852110655439
+0.7245342573849796,0.912855525440011,0.9447455998902117,0.96358115758054,0.9836346793838131,0.8073694350037202
+0.7228188149723814,0.9123752015644835,0.9446083644972039,0.9633753044910283,0.9835660616873092,0.8061516152582747
+0.7228016605482553,0.9123923559886095,0.945346004734621,0.9642501801214534,0.9842007753799705,0.8064601603251264
+0.7242426321748379,0.9127182900470031,0.9457577109136446,0.9643702610903352,0.9844924005901122,0.8075001427041348
+0.7249459635640032,0.9134730847085464,0.9458606374584005,0.9644217243627131,0.9844237828936082,0.8081538327780139
+0.7239338525405702,0.9126839811987512,0.9450200706762274,0.9633066867945244,0.9838405324733249,0.8070410918248734
+0.7236250729063025,0.9124266648368614,0.9448656808590936,0.9631866058256424,0.9836003705355612,0.806859810712907
+0.7224242632174838,0.911826259992452,0.9446941366178337,0.962723436374241,0.9832744364771675,0.8058241791370292
+0.722218410127972,0.9117404878718222,0.9443853569835661,0.962603355405359,0.9832401276289156,0.8056427934609479
+0.7226301163069956,0.9120149586578379,0.9448485264349675,0.9632552235221463,0.9833945174460493,0.8060240064701059
+0.7220297114625862,0.9115860980546883,0.9442138127423062,0.962689127525989,0.9828970391463958,0.805447713601981
+0.7237108450269324,0.9122722750197276,0.9448828352832196,0.9632723779462723,0.9833430541736714,0.8067396702769626
+0.7235736096339246,0.9124609736851134,0.9450715339486053,0.9634096133392802,0.9836346793838131,0.8067551081005684
+0.7245857206573575,0.913781864342814,0.9461007993961643,0.9646618863004769,0.9845781727107421,0.8077739135063374
+0.7238309259958143,0.9125124369574913,0.9448656808590936,0.9632895323703983,0.9833087453254195,0.8068440166640881
+0.7226815795793735,0.910934229937901,0.9437849521391567,0.962723436374241,0.9830171201152778,0.8056550566673708
+0.7228016605482553,0.9105911414553813,0.9439050331080385,0.9623803478917212,0.9831372010841596,0.8056326358407009
+0.7226129618828696,0.9121693484749717,0.9442652760146842,0.9631179881291385,0.9832915909012935,0.8061118839646912
+0.7240539335094521,0.9126153635022473,0.9447455998902117,0.9634096133392802,0.9834459807184273,0.8071669628005791
+0.7246200295056094,0.9130956873777747,0.9450886883727313,0.9632209146738944,0.9833773630219234,0.8075709724851876
+0.7247229560503654,0.9139190997358219,0.9458091741860226,0.9640786358801935,0.9842350842282225,0.8079037101451694
+0.7247058016262394,0.9128383710158849,0.9447970631625896,0.9633924589151542,0.9834802895666792,0.8075268394478943
+0.7238652348440663,0.9124609736851134,0.9447627543143376,0.9631351425532645,0.9833430541736714,0.8070234679633086
+0.7238652348440663,0.9124266648368614,0.9448828352832196,0.9633924589151542,0.9835489072631832,0.8069960834787101
+0.7244999485367276,0.9130785329536487,0.945397468006999,0.9638556283665557,0.9840635399869626,0.8075214826372252
+0.725117507805263,0.912889834288263,0.9453288503104951,0.9635983120046661,0.9838062236250729,0.8078535218173943
+0.7251518166535149,0.9133186948914125,0.9457405564895186,0.9640443270319415,0.9838919957457029,0.808052022558173
+0.725100353381137,0.9135931656774282,0.9460321816996603,0.9644045699385871,0.9844924005901122,0.8081660689818039
+0.7251346622293889,0.9133358493155385,0.9455347034000069,0.9641300991525714,0.9838405324733249,0.8080744792082037
+0.7249631179881292,0.9133701581637904,0.945397468006999,0.9638556283665557,0.9838233780491988,0.8079103037460894

models/nli.py

@@ -0,0 +1,267 @@
+import os
+import random
+import torch
+import math
+import argparse
+import transformers
+import logging
+
+from torch import nn
+from torch.nn import DataParallel
+from models.data_utils import build_batch, LoggingHandler, get_examples
+from tqdm import tqdm
+from sklearn.metrics import precision_recall_fscore_support
+from transformers import AutoConfig, RobertaModel, AutoModel, AutoTokenizer, AdamW
+
+import torch.nn.functional as F
+import numpy as np
+
+
+class NLI(nn.Module):
+    """
+    NLI model based on BERT (using the code from: https://github.com/yg211/bert_nli)
+    """
+
+    def __init__(self, model_path=None, device='cuda', parallel=False, debug=False, label_num=3, batch_size=16):
+        super(NLI, self).__init__()
+
+        lm = 'roberta-large'
+
+        if model_path is not None:
+            configuration = AutoConfig.from_pretrained(lm)
+            self.bert = RobertaModel(configuration)
+        else:
+            self.bert = AutoModel.from_pretrained(lm)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(lm)
+        self.vdim = 1024
+        self.max_length = 256
+
+        self.nli_head = nn.Linear(self.vdim, label_num)
+        self.batch_size = batch_size
+
+        if parallel:
+            self.bert = DataParallel(self.bert)
+
+        # load trained model
+        if model_path is not None:
+            sdict = torch.load(model_path, map_location=lambda storage, loc: storage)
+            self.load_state_dict(sdict, strict=False)
+
+        self.to(device)
+        self.device = device
+
+        self.debug = debug
+
+    def load_model(self, sdict):
+        if self.gpu:
+            self.load_state_dict(sdict)
+            self.to('cuda')
+        else:
+            self.load_state_dict(sdict)
+
+    def forward(self, sent_pair_list):
+        all_probs = None
+        iterator = range(0, len(sent_pair_list), self.batch_size)
+        if self.debug:
+            iterator = tqdm(iterator, desc='batch')
+        for batch_idx in iterator:
+            probs = self.ff(sent_pair_list[batch_idx:batch_idx + self.batch_size]).data.cpu().numpy()
+            if all_probs is None:
+                all_probs = probs
+            else:
+                all_probs = np.append(all_probs, probs, axis=0)
+        labels = []
+        for pp in all_probs:
+            ll = np.argmax(pp)
+            if ll == 0:
+                labels.append('entailment')
+            elif ll == 1:
+                labels.append('contradiction')
+            else:
+                labels.append('neutral')
+        return labels, all_probs
+
+    def ff(self, sent_pair_list):
+        ids, types, masks = build_batch(self.tokenizer, sent_pair_list, max_len=self.max_length)
+        if ids is None:
+            return None
+        ids_tensor = torch.tensor(ids)
+        #ypes_tensor = torch.tensor(types)
+        masks_tensor = torch.tensor(masks)
+
+        ids_tensor = ids_tensor.to(self.device)
+        #types_tensor = types_tensor.to(self.device)
+        masks_tensor = masks_tensor.to(self.device)
+        # self.bert.to('cuda')
+        # self.nli_head.to('cuda')
+
+        cls_vecs = self.bert(input_ids=ids_tensor, attention_mask=masks_tensor)[1]
+        logits = self.nli_head(cls_vecs)
+        predict_probs = F.log_softmax(logits, dim=1)
+        return predict_probs
+
+    def save(self, output_path, config_dic=None, acc=None):
+        if acc is None:
+            model_name = 'nli_large_2.state_dict'
+        else:
+            model_name = 'nli_large_2_acc{}.state_dict'.format(acc)
+        opath = os.path.join(output_path, model_name)
+        if config_dic is None:
+            torch.save(self.state_dict(), opath)
+        else:
+            torch.save(config_dic, opath)
+
+
+def get_scheduler(optimizer, scheduler: str, warmup_steps: int, t_total: int):
+    """
+    Returns the correct learning rate scheduler
+    """
+    scheduler = scheduler.lower()
+    if scheduler == 'constantlr':
+        return transformers.optimization.get_constant_schedule(optimizer)
+    elif scheduler == 'warmupconstant':
+        return transformers.optimization.get_constant_schedule_with_warmup(optimizer, warmup_steps)
+    elif scheduler == 'warmuplinear':
+        return transformers.optimization.get_linear_schedule_with_warmup(optimizer, warmup_steps, t_total)
+    elif scheduler == 'warmupcosine':
+        return transformers.optimization.get_cosine_schedule_with_warmup(optimizer, warmup_steps, t_total)
+    elif scheduler == 'warmupcosinewithhardrestarts':
+        return transformers.optimization.get_cosine_with_hard_restarts_schedule_with_warmup(optimizer, warmup_steps,
+                                                                                            t_total)
+    else:
+        raise ValueError("Unknown scheduler {}".format(scheduler))
+
+
+def train(model, optimizer, scheduler, train_data, dev_data, batch_size, fp16, gpu,
+          max_grad_norm, best_acc, model_save_path):
+    loss_fn = nn.CrossEntropyLoss()
+    model.train()
+
+    step_cnt = 0
+    for pointer in tqdm(range(0, len(train_data), batch_size), desc='training'):
+        step_cnt += 1
+        sent_pairs = []
+        labels = []
+        for i in range(pointer, pointer + batch_size):
+            if i >= len(train_data):
+                break
+            sents = train_data[i].get_texts()
+            sent_pairs.append(sents)
+            labels.append(train_data[i].get_label())
+        predicted_probs = model.ff(sent_pairs)
+        if predicted_probs is None:
+            continue
+        true_labels = torch.LongTensor(labels)
+        if gpu:
+            true_labels = true_labels.to('cuda')
+        loss = loss_fn(predicted_probs, true_labels)
+        if fp16:
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+            torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_grad_norm)
+        else:
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
+
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+
+        if step_cnt % 5000 == 0:
+            acc = evaluate(model, dev_data, mute=True)
+            logging.info('==> step {} dev acc: {}'.format(step_cnt, acc))
+            model.train()  # model was in eval mode in evaluate(); re-activate the train mode
+            if acc > best_acc:
+                best_acc = acc
+                logging.info('Saving model...')
+                model.save(model_save_path, model.state_dict())
+
+    return best_acc
+
+
+def parse_args():
+    ap = argparse.ArgumentParser("arguments for bert-nli training")
+    ap.add_argument('-b', '--batch_size', type=int, default=64, help='batch size')
+    ap.add_argument('-ep', '--epoch_num', type=int, default=10, help='epoch num')
+    ap.add_argument('--fp16', type=int, default=0, help='use apex mixed precision training (1) or not (0)')
+    ap.add_argument('--gpu', type=int, default=1, help='use gpu (1) or not (0)')
+    ap.add_argument('-ss', '--scheduler_setting', type=str, default='WarmupLinear',
+                    choices=['WarmupLinear', 'ConstantLR', 'WarmupConstant', 'WarmupCosine',
+                             'WarmupCosineWithHardRestarts'])
+    ap.add_argument('-mg', '--max_grad_norm', type=float, default=1., help='maximum gradient norm')
+    ap.add_argument('-wp', '--warmup_percent', type=float, default=0.1,
+                    help='how many percentage of steps are used for warmup')
+
+    args = ap.parse_args()
+    return args.batch_size, args.epoch_num, args.fp16, args.gpu, args.scheduler_setting, args.max_grad_norm, args.warmup_percent
+
+
+def evaluate(model, test_data, mute=False):
+    model.eval()
+    sent_pairs = [test_data[i].get_texts() for i in range(len(test_data))]
+    all_labels = [test_data[i].get_label() for i in range(len(test_data))]
+    _, probs = model(sent_pairs)
+    all_predict = [np.argmax(pp) for pp in probs]
+    assert len(all_predict) == len(all_labels)
+
+    acc = len([i for i in range(len(all_labels)) if all_predict[i] == all_labels[i]]) * 1. / len(all_labels)
+    prf = precision_recall_fscore_support(all_labels, all_predict, average=None, labels=[0, 1])
+
+    if not mute:
+        print('==>acc<==', acc)
+        print('==>precision-recall-f1<==\n', prf)
+
+    return acc
+
+
+if __name__ == '__main__':
+
+    batch_size, epoch_num, fp16, gpu, scheduler_setting, max_grad_norm, warmup_percent = parse_args()
+    fp16 = bool(fp16)
+    gpu = bool(gpu)
+
+    print('=====Arguments=====')
+    print('batch size:\t{}'.format(batch_size))
+    print('epoch num:\t{}'.format(epoch_num))
+    print('fp16:\t{}'.format(fp16))
+    print('gpu:\t{}'.format(gpu))
+    print('scheduler setting:\t{}'.format(scheduler_setting))
+    print('max grad norm:\t{}'.format(max_grad_norm))
+    print('warmup percent:\t{}'.format(warmup_percent))
+    print('=====Arguments=====')
+
+    label_num = 3
+    model_save_path = 'weights/entailment'
+
+    logging.basicConfig(format='%(asctime)s - %(message)s',
+                        datefmt='%Y-%m-%d %H:%M:%S',
+                        level=logging.INFO,
+                        handlers=[LoggingHandler()])
+
+    # Read the dataset
+    train_data = get_examples('../data/allnli/train.jsonl')
+    dev_data = get_examples('../data/allnli/dev.jsonl')
+
+    logging.info('train data size {}'.format(len(train_data)))
+    logging.info('dev data size {}'.format(len(dev_data)))
+    total_steps = math.ceil(epoch_num * len(train_data) * 1. / batch_size)
+    warmup_steps = int(total_steps * warmup_percent)
+
+    model = NLI(batch_size=batch_size, parallel=True)
+
+    optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-6, correct_bias=False)
+    scheduler = get_scheduler(optimizer, scheduler_setting, warmup_steps=warmup_steps, t_total=total_steps)
+    if fp16:
+        try:
+            from apex import amp
+        except ImportError:
+            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
+        model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
+
+    best_acc = -1.
+    for ep in range(epoch_num):
+        random.shuffle(train_data)
+        logging.info('\n=====epoch {}/{}====='.format(ep, epoch_num))
+        best_acc = train(model, optimizer, scheduler, train_data, dev_data, batch_size, fp16, gpu,
+                         max_grad_norm, best_acc, model_save_path)

models/qa_ranker.py

@@ -0,0 +1,284 @@
+import os
+import torch
+import logging
+import math
+import argparse
+import copy
+
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import transformers
+
+from torch import nn
+from torch.nn import DataParallel
+from transformers import BertModel, BertTokenizer
+
+from .data_utils import build_batch, LoggingHandler, get_examples, get_qa_examples
+
+from datetime import datetime
+from tqdm import tqdm
+from transformers import *
+from nltk.tokenize import word_tokenize
+from sklearn.metrics import precision_recall_fscore_support
+
+#os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+#os.environ["CUDA_VISIBLE_DEVICES"] = "8,15"
+
+
+class PassageRanker(nn.Module):
+    """Performs prediction, given the input of BERT embeddings.
+    """
+
+    def __init__(self, model_path=None, gpu=True, label_num=2, batch_size=16):
+        super(PassageRanker, self).__init__()
+
+        lm = 'bert-large-cased'
+
+        if model_path is not None:
+            configuration = AutoConfig.from_pretrained(lm)
+            self.language_model = BertModel(configuration)
+        else:
+            self.language_model = AutoModel.from_pretrained(lm)
+
+        self.language_model = DataParallel(self.language_model)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(lm)
+        self.vdim = 1024
+        self.max_length = 256
+
+        self.classification_head = nn.Linear(self.vdim, label_num)
+        self.gpu = gpu
+        self.batch_size = batch_size
+
+        # load trained model
+        if model_path is not None:
+            if gpu:
+                sdict = torch.load(model_path)
+                self.load_state_dict(sdict, strict=False)
+                self.to('cuda')
+            else:
+                sdict = torch.load(model_path, map_location=lambda storage, loc: storage)
+                self.load_state_dict(sdict, strict=False)
+        else:
+            if self.gpu:
+                self.to('cuda')
+
+    def load_model(self, sdict):
+        if self.gpu:
+            self.load_state_dict(sdict)
+            self.to('cuda')
+        else:
+            self.load_state_dict(sdict)
+
+    def forward(self, sent_pair_list):
+        all_probs = None
+        for batch_idx in range(0, len(sent_pair_list), self.batch_size):
+            probs = self.ff(sent_pair_list[batch_idx:batch_idx + self.batch_size]).data.cpu().numpy()
+            if all_probs is None:
+                all_probs = probs
+            else:
+                all_probs = np.append(all_probs, probs, axis=0)
+        labels = []
+        for pp in all_probs:
+            ll = np.argmax(pp)
+            if ll == 0:
+                labels.append('relevant')
+            else:
+                labels.append('irrelevant')
+        return labels, all_probs
+
+    def ff(self, sent_pair_list):
+        ids, types, masks = build_batch(self.tokenizer, sent_pair_list, max_len=self.max_length)
+        if ids is None:
+            return None
+        ids_tensor = torch.tensor(ids)
+        types_tensor = torch.tensor(types)
+        masks_tensor = torch.tensor(masks)
+
+        if self.gpu:
+            ids_tensor = ids_tensor.to('cuda')
+            types_tensor = types_tensor.to('cuda')
+            masks_tensor = masks_tensor.to('cuda')
+            # self.bert.to('cuda')
+            # self.nli_head.to('cuda')
+
+        cls_vecs = self.language_model(input_ids=ids_tensor, token_type_ids=types_tensor, attention_mask=masks_tensor)[1]
+        logits = self.classification_head(cls_vecs)
+        predict_probs = F.log_softmax(logits, dim=1)
+        return predict_probs
+
+    def save(self, output_path, config_dic=None, acc=None):
+        if acc is None:
+            model_name = 'qa_ranker.state_dict'
+        else:
+            model_name = 'qa_ranker_acc{}.state_dict'.format(acc)
+        opath = os.path.join(output_path, model_name)
+        if config_dic is None:
+            torch.save(self.state_dict(), opath)
+        else:
+            torch.save(config_dic, opath)
+
+    @staticmethod
+    def load(input_path, gpu=True, label_num=2, batch_size=16):
+        if gpu:
+            sdict = torch.load(input_path)
+        else:
+            sdict = torch.load(input_path, map_location=lambda storage, loc: storage)
+        model = PassageRanker(gpu=gpu, label_num=label_num, batch_size=batch_size)
+        model.load_state_dict(sdict)
+        return model
+
+
+def get_scheduler(optimizer, scheduler: str, warmup_steps: int, t_total: int):
+    """
+    Returns the correct learning rate scheduler
+    """
+    scheduler = scheduler.lower()
+    if scheduler == 'constantlr':
+        return transformers.optimization.get_constant_schedule(optimizer)
+    elif scheduler == 'warmupconstant':
+        return transformers.optimization.get_constant_schedule_with_warmup(optimizer, warmup_steps)
+    elif scheduler == 'warmuplinear':
+        return transformers.optimization.get_linear_schedule_with_warmup(optimizer, warmup_steps, t_total)
+    elif scheduler == 'warmupcosine':
+        return transformers.optimization.get_cosine_schedule_with_warmup(optimizer, warmup_steps, t_total)
+    elif scheduler == 'warmupcosinewithhardrestarts':
+        return transformers.optimization.get_cosine_with_hard_restarts_schedule_with_warmup(optimizer, warmup_steps,
+                                                                                            t_total)
+    else:
+        raise ValueError("Unknown scheduler {}".format(scheduler))
+
+
+def train(model, optimizer, scheduler, train_data, dev_data, batch_size, fp16, gpu,
+          max_grad_norm, best_acc, model_save_path):
+    loss_fn = nn.CrossEntropyLoss()
+    model.train()
+
+    step_cnt = 0
+    for pointer in tqdm(range(0, len(train_data), batch_size), desc='training'):
+        step_cnt += 1
+        sent_pairs = []
+        labels = []
+        for i in range(pointer, pointer + batch_size):
+            if i >= len(train_data):
+                break
+            sents = train_data[i].get_texts()
+            sent_pairs.append(sents)
+            labels.append(train_data[i].get_label())
+        predicted_probs = model.ff(sent_pairs)
+        if predicted_probs is None:
+            continue
+        true_labels = torch.LongTensor(labels)
+        if gpu:
+            true_labels = true_labels.to('cuda')
+        loss = loss_fn(predicted_probs, true_labels)
+        if fp16:
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+            torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_grad_norm)
+        else:
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
+
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+
+        if step_cnt % 5000 == 0:
+            acc = evaluate(model, dev_data, mute=True)
+            logging.info('==> step {} dev acc: {}'.format(step_cnt, acc))
+            model.train()  # model was in eval mode in evaluate(); re-activate the train mode
+            if acc > best_acc:
+                best_acc = acc
+                logging.info('Saving model...')
+                model.save(model_save_path, model.state_dict())
+
+    return best_acc
+
+
+def parse_args():
+    ap = argparse.ArgumentParser("arguments for bert-nli training")
+    ap.add_argument('-b', '--batch_size', type=int, default=128, help='batch size')
+    ap.add_argument('-ep', '--epoch_num', type=int, default=10, help='epoch num')
+    ap.add_argument('--fp16', type=int, default=0, help='use apex mixed precision training (1) or not (0)')
+    ap.add_argument('--gpu', type=int, default=1, help='use gpu (1) or not (0)')
+    ap.add_argument('-ss', '--scheduler_setting', type=str, default='WarmupLinear',
+                    choices=['WarmupLinear', 'ConstantLR', 'WarmupConstant', 'WarmupCosine',
+                             'WarmupCosineWithHardRestarts'])
+    ap.add_argument('-mg', '--max_grad_norm', type=float, default=1., help='maximum gradient norm')
+    ap.add_argument('-wp', '--warmup_percent', type=float, default=0.1,
+                    help='how many percentage of steps are used for warmup')
+
+    args = ap.parse_args()
+    return args.batch_size, args.epoch_num, args.fp16, args.gpu, args.scheduler_setting, args.max_grad_norm, args.warmup_percent
+
+
+def evaluate(model, test_data, mute=False):
+    model.eval()
+    sent_pairs = [test_data[i].get_texts() for i in range(len(test_data))]
+    all_labels = [test_data[i].get_label() for i in range(len(test_data))]
+    _, probs = model(sent_pairs)
+    all_predict = [np.argmax(pp) for pp in probs]
+    assert len(all_predict) == len(all_labels)
+
+    acc = len([i for i in range(len(all_labels)) if all_predict[i] == all_labels[i]]) * 1. / len(all_labels)
+    prf = precision_recall_fscore_support(all_labels, all_predict, average=None, labels=[0, 1])
+
+    if not mute:
+        print('==>acc<==', acc)
+        print('label meanings: 0: relevant, 1: irrelevant')
+        print('==>precision-recall-f1<==\n', prf)
+
+    return acc
+
+
+if __name__ == '__main__':
+
+    batch_size, epoch_num, fp16, gpu, scheduler_setting, max_grad_norm, warmup_percent = parse_args()
+    fp16 = bool(fp16)
+    gpu = bool(gpu)
+
+    print('=====Arguments=====')
+    print('batch size:\t{}'.format(batch_size))
+    print('epoch num:\t{}'.format(epoch_num))
+    print('fp16:\t{}'.format(fp16))
+    print('gpu:\t{}'.format(gpu))
+    print('scheduler setting:\t{}'.format(scheduler_setting))
+    print('max grad norm:\t{}'.format(max_grad_norm))
+    print('warmup percent:\t{}'.format(warmup_percent))
+    print('=====Arguments=====')
+
+    label_num = 2
+    model_save_path = 'weights/passage_ranker_2'
+
+    logging.basicConfig(format='%(asctime)s - %(message)s',
+                        datefmt='%Y-%m-%d %H:%M:%S',
+                        level=logging.INFO,
+                        handlers=[LoggingHandler()])
+
+    # Read the dataset
+    train_data = get_qa_examples('../data/qa_ranking_large/train.jsonl', dev=False)
+    dev_data = get_qa_examples('../data/qa_ranking_large/dev.jsonl', dev=True)[:20000]
+
+    logging.info('train data size {}'.format(len(train_data)))
+    logging.info('dev data size {}'.format(len(dev_data)))
+    total_steps = math.ceil(epoch_num * len(train_data) * 1. / batch_size)
+    warmup_steps = int(total_steps * warmup_percent)
+
+    model = PassageRanker(gpu=gpu, batch_size=batch_size)
+
+    optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-6, correct_bias=False)
+    scheduler = get_scheduler(optimizer, scheduler_setting, warmup_steps=warmup_steps, t_total=total_steps)
+    if fp16:
+        try:
+            from apex import amp
+        except ImportError:
+            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
+        model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
+
+    best_acc = -1.
+    for ep in range(epoch_num):
+        logging.info('\n=====epoch {}/{}====='.format(ep, epoch_num))
+        best_acc = train(model, optimizer, scheduler, train_data, dev_data, batch_size, fp16, gpu,
+                         max_grad_norm, best_acc, model_save_path)

models/sparse_retriever.py

@@ -0,0 +1,204 @@
+import hashlib
+import logging
+import math
+import re
+
+from multiprocessing import Pool
+from nltk import WordNetLemmatizer, pos_tag
+from nltk.corpus import wordnet, stopwords
+from typing import List
+from tqdm import tqdm
+
+lemmatizer = WordNetLemmatizer()
+stopwords = set(stopwords.words('english'))
+
+
+def get_ngrams(text_tokens: List[str], min_length=1, max_length=4) -> List[str]:
+    """
+    Gets word-level ngrams from text
+    :param text_tokens: the string used to generate ngrams
+    :param min_length: the minimum length og the generated ngrams in words
+    :param max_length: the maximum length og the generated ngrams in words
+    :return: list of ngrams (strings)
+    """
+    max_length = min(max_length, len(text_tokens))
+    all_ngrams = []
+    for n in range(min_length - 1, max_length):
+        ngrams = [" ".join(ngram) for ngram in zip(*[text_tokens[i:] for i in range(n + 1)])]
+        for ngram in ngrams:
+            if '#' not in ngram:
+                all_ngrams.append(ngram)
+
+    return all_ngrams
+
+
+def clean_text(text):
+    """
+    Removes non-alphanumeric symbols from text
+    :param text:
+    :return: clean text
+    """
+    text = text.replace('-', ' ')
+    text = re.sub('[^a-zA-Z0-9 ]+', '', text)
+    text = re.sub(' +', ' ', text)
+    return text
+
+
+def string_hash(string):
+    """
+    Returns a static hash value for a string
+    :param string:
+    :return:
+    """
+    return int(hashlib.md5(string.encode('utf8')).hexdigest(), 16)
+
+
+def tokenize(text, lemmatize=True, ngrams_length=2):
+    """
+    :param text:
+    :param stopwords: set of stopwords to exclude
+    :param lemmatize:
+    :param ngrams_length: the maximum number of tokens per ngram
+    :return:
+    """
+    tokens = clean_text(text).lower().split(' ')
+    tokens = [t for t in tokens if t != '']
+    if lemmatize:
+        lemmatized_words = []
+        pos_labels = pos_tag(tokens)
+        pos_labels = [pos[1][0].lower() for pos in pos_labels]
+
+        for i, word in enumerate(tokens):
+            if word in stopwords:
+                word = '#'
+            if pos_labels[i] == 'j':
+                pos_labels[i] = 'a'  # 'j' <--> 'a' reassignment
+            if pos_labels[i] in ['r']:  # For adverbs it's a bit different
+                try:
+                    lemma = wordnet.synset(word + '.r.1').lemmas()[0].pertainyms()[0].name()
+                except:
+                    lemma = word
+                lemmatized_words.append(lemma)
+            elif pos_labels[i] in ['a', 's', 'v']:  # For adjectives and verbs
+                lemmatized_words.append(lemmatizer.lemmatize(word, pos=pos_labels[i]))
+            else:  # For nouns and everything else as it is the default kwarg
+                lemmatized_words.append(lemmatizer.lemmatize(word))
+        tokens = lemmatized_words
+
+    ngrams = get_ngrams(tokens, max_length=ngrams_length)
+
+    return ngrams
+
+
+class SparseRetriever:
+    def __init__(self, ngram_buckets=16777216, k1=1.5, b=0.75, epsilon=0.25,
+                 max_relative_freq=0.5, workers=4):
+        self.ngram_buckets = ngram_buckets
+        self.k1 = k1
+        self.b = b
+        self.epsilon = epsilon
+        self.max_relative_freq = max_relative_freq
+        self.corpus_size = 0
+        self.avgdl = 0
+
+        self.inverted_index = {}
+        self.idf = {}
+        self.doc_len = []
+        self.workers = workers
+
+    def index_documents(self, documents):
+        with Pool(self.workers) as p:
+            tokenized_documents = list(tqdm(p.imap(tokenize, documents), total=len(documents), desc='tokenized'))
+
+        logging.info('Building inverted index...')
+
+        self.corpus_size = len(tokenized_documents)
+        nd = self._create_inverted_index(tokenized_documents)
+        self._calc_idf(nd)
+
+        logging.info('Built inverted index')
+
+    def _create_inverted_index(self, documents):
+        nd = {}  # word -> number of documents with word
+        num_doc = 0
+        for doc_id, document in enumerate(tqdm(documents, desc='indexed')):
+            self.doc_len.append(len(document))
+            num_doc += len(document)
+
+            frequencies = {}
+            for word in document:
+                if word not in frequencies:
+                    frequencies[word] = 0
+                frequencies[word] += 1
+
+            for word, freq in frequencies.items():
+                hashed_word = string_hash(word) % self.ngram_buckets
+                if hashed_word not in self.inverted_index:
+                    self.inverted_index[hashed_word] = [(doc_id, freq)]
+                else:
+                    self.inverted_index[hashed_word].append((doc_id, freq))
+
+            for word, freq in frequencies.items():
+                hashed_word = string_hash(word) % self.ngram_buckets
+                if hashed_word not in nd:
+                    nd[hashed_word] = 0
+                nd[hashed_word] += 1
+
+        self.avgdl = num_doc / self.corpus_size
+
+        return nd
+
+    def _calc_idf(self, nd):
+        """
+        Calculates frequencies of terms in documents and in corpus.
+        This algorithm sets a floor on the idf values to eps * average_idf
+        """
+        # collect idf sum to calculate an average idf for epsilon value
+        idf_sum = 0
+        # collect words with negative idf to set them a special epsilon value.
+        # idf can be negative if word is contained in more than half of documents
+        negative_idfs = []
+        for word, freq in nd.items():
+            idf = math.log(self.corpus_size - freq + 0.5) - math.log(freq + 0.5)
+            r_freq = float(freq) / self.corpus_size
+            if r_freq > self.max_relative_freq:
+                continue
+            self.idf[word] = idf
+            idf_sum += idf
+            if idf < 0:
+                negative_idfs.append(word)
+        self.average_idf = idf_sum / len(self.idf)
+
+        eps = self.epsilon * self.average_idf
+        for word in negative_idfs:
+            self.idf[word] = eps
+
+    def _get_scores(self, query):
+        """
+        The ATIRE BM25 variant uses an idf function which uses a log(idf) score. To prevent negative idf scores,
+        this algorithm also adds a floor to the idf value of epsilon.
+        See [Trotman, A., X. Jia, M. Crane, Towards an Efficient and Effective Search Engine] for more info
+        :param query:
+        :return:
+        """
+        query = tokenize(query)
+        scores = {}
+        for q in query:
+            hashed_word = string_hash(q) % self.ngram_buckets
+            idf = self.idf.get(hashed_word)
+            if idf:
+                doc_freqs = self.inverted_index[hashed_word]
+                for doc_id, freq in doc_freqs:
+                    score = idf * (freq * (self.k1 + 1) /
+                                   (freq + self.k1 * 1 - self.b + (self.b * self.doc_len[doc_id] / self.avgdl)))
+                    if doc_id in scores:
+                        scores[doc_id] += score
+                    else:
+                        scores[doc_id] = score
+        return sorted(scores.items(), key=lambda x: x[1], reverse=True)
+
+    def search(self, queries, topk=100):
+        results = [self._get_scores(q) for q in tqdm(queries, desc='searched')]
+        results = [r[:topk] for r in results]
+        logging.info('Done searching')
+        return results

models/sparse_retriever_fast.py

@@ -0,0 +1,48 @@
+import logging
+import os
+import tantivy
+
+from tqdm import tqdm
+
+
+class SparseRetrieverFast:
+    def __init__(self, path='sparse_index', load=True):
+        if not os.path.exists(path):
+            os.mkdir(path)
+        schema_builder = tantivy.SchemaBuilder()
+        schema_builder.add_text_field("body", stored=False)
+        schema_builder.add_unsigned_field("doc_id", stored=True)
+        schema = schema_builder.build()
+        self.index = tantivy.Index(schema, path=path, reuse=load)
+        self.searcher = self.index.searcher()
+
+    def index_documents(self, documents):
+        logging.info('Building sparse index of {} docs...'.format(len(documents)))
+        writer = self.index.writer()
+        for i, doc in enumerate(documents):
+            writer.add_document(tantivy.Document(
+                body=[doc],
+                doc_id=i
+            ))
+            if (i+1) % 100000 == 0:
+                writer.commit()
+                logging.info('Indexed {} docs'.format(i+1))
+        writer.commit()
+        logging.info('Built sparse index')
+        self.index.reload()
+        self.searcher = self.index.searcher()
+
+    def search(self, queries, topk=100):
+        results = []
+        for q in tqdm(queries, desc='searched'):
+            docs = []
+            try:
+                query = self.index.parse_query(q, ["body"])
+                scores = self.searcher.search(query, topk).hits
+                docs = [(self.searcher.doc(doc_id)['doc_id'][0], score)
+                        for score, doc_id in scores]
+            except:
+                pass
+            results.append(docs)
+
+        return results

models/text_encoder.py

@@ -0,0 +1,867 @@
+import importlib
+import json
+import logging
+import os
+
+import hashlib
+import pickle
+import torch
+torch.device(0)
+import transformers
+
+import numpy as np
+
+from pprint import pprint
+from typing import Iterable, Tuple, Type, Dict, List, Union, Optional
+from numpy.core.multiarray import ndarray
+from torch import nn
+from torch.nn import DataParallel
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader, Dataset
+
+from models.data_utils import get_qnli_examples, get_single_examples, get_ict_examples, get_examples, get_qar_examples, \
+    get_qar_artificial_examples, get_retrieval_examples
+from tqdm import tqdm
+
+from collections import OrderedDict
+from transformers import BertTokenizer, BertModel, AutoTokenizer, AutoModel
+
+from models.vector_index import VectorIndex
+
+
+class InputExample:
+    """
+    Structure for one input example with texts, the label and a unique id
+    """
+    def __init__(self, guid: str, texts: List[str], label: Union[int, float]):
+        """
+        Creates one InputExample with the given texts, guid and label
+
+        str.strip() is called on both texts.
+
+        :param guid
+            id for the example
+        :param texts
+            the texts for the example
+        :param label
+            the label for the example
+        """
+        self.guid = guid
+        self.texts = [text.strip() for text in texts]
+        self.label = label
+
+
+class LoggingHandler(logging.Handler):
+    def __init__(self, level=logging.NOTSET):
+        super().__init__(level)
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            tqdm.write(msg)
+            self.flush()
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except:
+            self.handleError(record)
+
+
+def import_from_string(dotted_path):
+    """
+    Import a dotted module path and return the attribute/class designated by the
+    last name in the path. Raise ImportError if the import failed.
+    """
+    try:
+        module_path, class_name = dotted_path.rsplit('.', 1)
+    except ValueError:
+        msg = "%s doesn't look like a module path" % dotted_path
+        raise ImportError(msg)
+
+    module = importlib.import_module(module_path)
+
+    try:
+        return getattr(module, class_name)
+    except AttributeError:
+        msg = 'Module "%s" does not define a "%s" attribute/class' % (module_path, class_name)
+        raise ImportError(msg)
+
+
+def batch_to_device(batch, target_device: torch.device):
+    """
+    send a batch to a device
+
+    :param batch:
+    :param target_device:
+    :return: the batch sent to the device
+    """
+    features = batch['features']
+    for paired_sentence_idx in range(len(features)):
+        for feature_name in features[paired_sentence_idx]:
+            features[paired_sentence_idx][feature_name] = features[paired_sentence_idx][feature_name].to(target_device)
+
+    labels = batch['labels'].to(target_device)
+    return features, labels
+
+
+class BERT(nn.Module):
+    """BERT model to generate token embeddings.
+    Each token is mapped to an output vector from BERT.
+    """
+    def __init__(self, model_name_or_path: str, max_seq_length: int = 128, do_lower_case: Optional[bool] = None, model_args: Dict = {}, tokenizer_args: Dict = {}):
+        super(BERT, self).__init__()
+        self.config_keys = ['max_seq_length', 'do_lower_case']
+        self.do_lower_case = do_lower_case
+
+        if max_seq_length > 510:
+            logging.warning("BERT only allows a max_seq_length of 510 (512 with special tokens). Value will be set to 510")
+            max_seq_length = 510
+        self.max_seq_length = max_seq_length
+
+        if self.do_lower_case is not None:
+            tokenizer_args['do_lower_case'] = do_lower_case
+
+        self.model = AutoModel.from_pretrained(model_name_or_path, **model_args)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, **tokenizer_args)
+
+    def forward(self, features):
+        """Returns token_embeddings, cls_token"""
+        output_states = self.model(**features)
+        output_tokens = output_states[0]
+        cls_tokens = output_tokens[:, 0, :]  # CLS token is first token
+        features.update({'token_embeddings': output_tokens, 'cls_token_embeddings': cls_tokens, 'attention_mask': features['attention_mask']})
+
+        if len(output_states) > 2:
+            features.update({'all_layer_embeddings': output_states[2]})
+
+        return features
+
+    def get_word_embedding_dimension(self) -> int:
+        return self.model.config.hidden_size
+
+    def tokenize(self, text: str) -> List[int]:
+        """
+        Tokenizes a text and maps tokens to token-ids
+        """
+        return self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(text))
+
+    def get_sentence_features(self, tokens: List[int], pad_seq_length: int):
+        """
+        Convert tokenized sentence in its embedding ids, segment ids and mask
+        :param tokens:
+            a tokenized sentence
+        :param pad_seq_length:
+            the maximal length of the sequence. Cannot be greater than self.sentence_transformer_config.max_seq_length
+        :return: embedding ids, segment ids and mask for the sentence
+        """
+        pad_seq_length = min(pad_seq_length, self.max_seq_length) + 2  ##Add Space for CLS + SEP token
+
+        return self.tokenizer.prepare_for_model(tokens, max_length=pad_seq_length, pad_to_max_length=True, return_tensors='pt')
+
+    def get_config_dict(self):
+        return {key: self.__dict__[key] for key in self.config_keys}
+
+    def save(self, output_path: str):
+        self.model.save_pretrained(output_path)
+        self.tokenizer.save_pretrained(output_path)
+
+        with open(os.path.join(output_path, 'sentence_bert_config.json'), 'w') as fOut:
+            json.dump(self.get_config_dict(), fOut, indent=2)
+
+    @staticmethod
+    def load(input_path: str):
+        with open(os.path.join(input_path, 'sentence_bert_config.json')) as fIn:
+            config = json.load(fIn)
+        return BERT(model_name_or_path=input_path, **config)
+
+
+class SentenceTransformer(nn.Sequential):
+    def __init__(self, model_path: str = None, modules: Iterable[nn.Module] = None, device: str = None,
+                 parallel=False):
+        if modules is not None and not isinstance(modules, OrderedDict):
+            modules = OrderedDict([(str(idx), module) for idx, module in enumerate(modules)])
+
+        if model_path is not None:
+            logging.info("Load SentenceTransformer from folder: {}".format(model_path))
+
+            with open(os.path.join(model_path, 'modules.json')) as fIn:
+                contained_modules = json.load(fIn)
+
+            modules = OrderedDict()
+            for module_config in contained_modules:
+                module_class = import_from_string(module_config['type'])
+                module = module_class.load(os.path.join(model_path, module_config['path']))
+                if 'BERT' in module_config['type']:
+                    if parallel:
+                        module = DataParallel(module)
+                modules[module_config['name']] = module
+
+        super().__init__(modules)
+
+        self.best_score = -1
+        self.total_steps = 0
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            logging.info("Use pytorch device: {}".format(device))
+        self.device = torch.device(device)
+        self.to(device)
+
+    def _eval_during_training(self, evaluator, output_path, save_best_model, epoch, steps):
+        """Runs evaluation during the training"""
+        if evaluator is not None:
+            score = evaluator(self, output_path=output_path, epoch=epoch, steps=steps)
+            print(score)
+            if score > self.best_score and save_best_model:
+                print('saving')
+                self.save(output_path)
+                self.best_score = score
+
+    def evaluate(self, evaluator, output_path: str = None):
+        """
+        Evaluate the model
+
+        :param evaluator:
+            the evaluator
+        :param output_path:
+            the evaluator can write the results to this path
+        """
+        if output_path is not None:
+            os.makedirs(output_path, exist_ok=True)
+        return evaluator(self, output_path)
+
+    def encode(self, sentences: List[str], batch_size: int = 8, show_progress_bar: bool = None, output_value: str = 'sentence_embedding', convert_to_numpy: bool = True) -> List[ndarray]:
+        """
+        Computes sentence embeddings
+        :param sentences:
+           the sentences to embed
+        :param batch_size:
+           the batch size used for the computation
+        :param show_progress_bar:
+            Output a progress bar when encode sentences
+        :param output_value:
+            Default sentence_embedding, to get sentence embeddings. Can be set to token_embeddings
+            to get wordpiece token embeddings.
+        :param convert_to_numpy:
+            If true, the output is a list of numpy vectors. Else, it is a list of pytorch tensors.
+        :return:
+           Depending on convert_to_numpy, either a list of numpy vectors or a list of pytorch tensors
+        """
+        self.eval()
+        if show_progress_bar is None:
+            show_progress_bar = (logging.getLogger().getEffectiveLevel()==logging.INFO or logging.getLogger().getEffectiveLevel()==logging.DEBUG)
+
+        all_embeddings = []
+        length_sorted_idx = np.argsort([len(sen) for sen in sentences])
+
+        iterator = range(0, len(sentences), batch_size)
+        if show_progress_bar:
+            iterator = tqdm(iterator, desc="Batches")
+
+        for batch_idx in iterator:
+            batch_tokens = []
+
+            batch_start = batch_idx
+            batch_end = min(batch_start + batch_size, len(sentences))
+
+            longest_seq = 0
+
+            for idx in length_sorted_idx[batch_start: batch_end]:
+                sentence = sentences[idx]
+                tokens = self.tokenize(sentence)
+                longest_seq = max(longest_seq, len(tokens))
+                batch_tokens.append(tokens)
+
+            features = {}
+            for text in batch_tokens:
+                sentence_features = self.get_sentence_features(text, longest_seq)
+
+                for feature_name in sentence_features:
+                    if feature_name not in features:
+                        features[feature_name] = []
+                    features[feature_name].append(sentence_features[feature_name])
+
+            for feature_name in features:
+                try:
+                    features[feature_name] = torch.tensor(np.asarray(features[feature_name])).to(self.device)
+                except:
+                    features[feature_name] = torch.cat(features[feature_name]).to(self.device)
+
+            with torch.no_grad():
+                out_features = self.forward(features)
+                embeddings = out_features[output_value]
+
+                if output_value == 'token_embeddings':
+                    #Set token embeddings to 0 for padding tokens
+                    input_mask = out_features['attention_mask']
+                    input_mask_expanded = input_mask.unsqueeze(-1).expand(embeddings.size()).float()
+                    embeddings = embeddings * input_mask_expanded
+
+                if convert_to_numpy:
+                    embeddings = embeddings.to('cpu').numpy()
+
+                all_embeddings.extend(embeddings)
+
+        reverting_order = np.argsort(length_sorted_idx)
+        all_embeddings = [all_embeddings[idx] for idx in reverting_order]
+
+        return all_embeddings
+
+    def get_max_seq_length(self):
+        if hasattr(self._first_module(), 'max_seq_length'):
+            return self._first_module().max_seq_length
+
+        return None
+
+    def tokenize(self, text):
+        return self._first_module().tokenize(text)
+
+    def get_sentence_features(self, *features):
+        return self._first_module().get_sentence_features(*features)
+
+    def get_sentence_embedding_dimension(self):
+        return self._last_module().get_sentence_embedding_dimension()
+
+    def _first_module(self):
+        """Returns the first module of this sequential embedder"""
+        try:
+            return self._modules[next(iter(self._modules))].module
+        except:
+            return self._modules[next(iter(self._modules))]
+
+    def _last_module(self):
+        """Returns the last module of this sequential embedder"""
+        return self._modules[next(reversed(self._modules))]
+
+    def save(self, path):
+        """
+        Saves all elements for this seq. sentence embedder into different sub-folders
+        """
+        if path is None:
+            return
+
+        logging.info("Save model to {}".format(path))
+        contained_modules = []
+
+        for idx, name in enumerate(self._modules):
+            module = self._modules[name]
+            if isinstance(module, DataParallel):
+                module = module.module
+            model_path = os.path.join(path, str(idx) + "_" + type(module).__name__)
+            os.makedirs(model_path, exist_ok=True)
+            module.save(model_path)
+            contained_modules.append(
+                {'idx': idx, 'name': name, 'path': os.path.basename(model_path), 'type': type(module).__module__})
+
+        with open(os.path.join(path, 'modules.json'), 'w') as fOut:
+            json.dump(contained_modules, fOut, indent=2)
+
+        with open(os.path.join(path, 'config.json'), 'w') as fOut:
+            json.dump({'__version__': '1.0'}, fOut, indent=2)
+
+    def smart_batching_collate(self, batch):
+        """
+        Transforms a batch from a SmartBatchingDataset to a batch of tensors for the model
+        :param batch:
+            a batch from a SmartBatchingDataset
+        :return:
+            a batch of tensors for the model
+        """
+        num_texts = len(batch[0][0])
+
+        labels = []
+        paired_texts = [[] for _ in range(num_texts)]
+        max_seq_len = [0] * num_texts
+        for tokens, label in batch:
+            labels.append(label)
+            for i in range(num_texts):
+                paired_texts[i].append(tokens[i])
+                max_seq_len[i] = max(max_seq_len[i], len(tokens[i]))
+
+        features = []
+        for idx in range(num_texts):
+            max_len = max_seq_len[idx]
+            feature_lists = {}
+
+            for text in paired_texts[idx]:
+                sentence_features = self.get_sentence_features(text, max_len)
+
+                for feature_name in sentence_features:
+                    if feature_name not in feature_lists:
+                        feature_lists[feature_name] = []
+
+                    feature_lists[feature_name].append(sentence_features[feature_name])
+
+            for feature_name in feature_lists:
+                try:
+                    feature_lists[feature_name] = torch.tensor(np.asarray(feature_lists[feature_name]))
+                except:
+                    feature_lists[feature_name] = torch.cat(feature_lists[feature_name])
+
+            features.append(feature_lists)
+
+        return {'features': features, 'labels': torch.stack(labels)}
+
+    def fit(self,
+            train_objectives: Iterable[Tuple[DataLoader, nn.Module]],
+            evaluator,
+            epochs: int = 1,
+            steps_per_epoch=None,
+            scheduler: str = 'WarmupLinear',
+            warmup_steps: int = 10000,
+            optimizer_class: Type[Optimizer] = transformers.AdamW,
+            optimizer_params: Dict[str, object] = {'lr': 2e-5, 'eps': 1e-6, 'correct_bias': False},
+            acc_steps: int = 4,
+            weight_decay: float = 0.01,
+            evaluation_steps: int = 0,
+            output_path: str = None,
+            save_best_model: bool = True,
+            max_grad_norm: float = 1,
+            fp16: bool = False,
+            fp16_opt_level: str = 'O1',
+            local_rank: int = -1
+            ):
+        if output_path is not None:
+            os.makedirs(output_path, exist_ok=True)
+
+        dataloaders = [dataloader for dataloader, _ in train_objectives]
+
+        # Use smart batching
+        for dataloader in dataloaders:
+            dataloader.collate_fn = self.smart_batching_collate
+
+        loss_models = [loss for _, loss in train_objectives]
+        device = self.device
+
+        for loss_model in loss_models:
+            loss_model.to(device)
+
+        if steps_per_epoch is None or steps_per_epoch == 0:
+            steps_per_epoch = min([len(dataloader) for dataloader in dataloaders])
+
+        num_train_steps = int(steps_per_epoch * epochs)
+
+        # Prepare optimizers
+        optimizers = []
+        schedulers = []
+        for loss_model in loss_models:
+            param_optimizer = list(loss_model.named_parameters())
+
+            no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+            optimizer_grouped_parameters = [
+                {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
+                 'weight_decay': weight_decay},
+                {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+            ]
+            t_total = num_train_steps
+            if local_rank != -1:
+                t_total = t_total // torch.distributed.get_world_size()
+
+            optimizer = optimizer_class(optimizer_grouped_parameters, **optimizer_params)
+            scheduler_obj = self._get_scheduler(optimizer, scheduler=scheduler, warmup_steps=warmup_steps,
+                                                t_total=t_total)
+
+            optimizers.append(optimizer)
+            schedulers.append(scheduler_obj)
+
+        if fp16:
+            try:
+                from apex import amp
+            except ImportError:
+                raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
+
+            for train_idx in range(len(loss_models)):
+                model, optimizer = amp.initialize(loss_models[train_idx], optimizers[train_idx],
+                                                  opt_level=fp16_opt_level)
+                loss_models[train_idx] = model
+                optimizers[train_idx] = optimizer
+
+        data_iterators = [iter(dataloader) for dataloader in dataloaders]
+        num_train_objectives = len(train_objectives)
+
+        for epoch in range(epochs):
+            #logging.info('Epoch {}'.format(epoch))
+            epoch_loss = 0
+            epoch_steps = 0
+            for loss_model in loss_models:
+                loss_model.zero_grad()
+                loss_model.train()
+
+            step_iterator = tqdm(range(steps_per_epoch), desc='loss: -')
+            for step in step_iterator:
+                for train_idx in range(num_train_objectives):
+                    loss_model = loss_models[train_idx]
+                    optimizer = optimizers[train_idx]
+                    scheduler = schedulers[train_idx]
+                    data_iterator = data_iterators[train_idx]
+
+                    try:
+                        data = next(data_iterator)
+                    except StopIteration:
+                        # logging.info("Restart data_iterator")
+                        data_iterator = iter(dataloaders[train_idx])
+                        data_iterators[train_idx] = data_iterator
+                        data = next(data_iterator)
+
+                    features, labels = batch_to_device(data, self.device)
+                    loss_value = loss_model(features, labels)
+                    loss_value = loss_value / acc_steps
+
+                    if fp16:
+                        with amp.scale_loss(loss_value, optimizer) as scaled_loss:
+                            scaled_loss.backward()
+                        torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_grad_norm)
+                    else:
+                        loss_value.backward()
+                        torch.nn.utils.clip_grad_norm_(loss_model.parameters(), max_grad_norm)
+
+                    if (step + 1) % acc_steps == 0:
+                        optimizer.step()
+                        scheduler.step()
+                        optimizer.zero_grad()
+
+                self.total_steps += 1
+
+                if (step + 1) % acc_steps == 0:
+                    epoch_steps += 1
+                    epoch_loss += loss_value.item()
+                    step_iterator.set_description('loss: {} - acc steps: {}'.format((epoch_loss / epoch_steps),
+                                                                                    (self.total_steps / acc_steps)))
+
+                if evaluation_steps > 0 and self.total_steps > 0 and (self.total_steps / acc_steps) % evaluation_steps == 0:
+                    self._eval_during_training(evaluator, output_path, save_best_model, epoch, epoch_steps)
+                    for loss_model in loss_models:
+                        loss_model.zero_grad()
+                        loss_model.train()
+
+            self._eval_during_training(evaluator, output_path, save_best_model, epoch, -1)
+
+    def _get_scheduler(self, optimizer, scheduler: str, warmup_steps: int, t_total: int):
+        """
+        Returns the correct learning rate scheduler
+        """
+        scheduler = scheduler.lower()
+        if scheduler == 'constantlr':
+            return transformers.get_constant_schedule(optimizer)
+        elif scheduler == 'warmupconstant':
+            return transformers.get_constant_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps)
+        elif scheduler == 'warmuplinear':
+            return transformers.get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total)
+        elif scheduler == 'warmupcosine':
+            return transformers.get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total)
+        elif scheduler == 'warmupcosinewithhardrestarts':
+            return transformers.get_cosine_with_hard_restarts_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total)
+        else:
+            raise ValueError("Unknown scheduler {}".format(scheduler))
+
+
+class SentencesDataset(Dataset):
+    def __init__(self, examples: List[InputExample], model: SentenceTransformer, show_progress_bar: bool = None):
+        self.examples = examples
+        self.model = model
+
+    def __getitem__(self, item):
+        tokenized_texts = [model.tokenize(text) for text in self.examples[item].texts]
+        return tokenized_texts, torch.tensor(self.examples[item].label, dtype=torch.float)
+
+    def __len__(self):
+        return len(self.examples)
+
+
+class MultipleNegativesRankingLossANN(nn.Module):
+    def __init__(self, model: SentenceTransformer, negative_samples=4):
+        super(MultipleNegativesRankingLossANN, self).__init__()
+        self.model = model
+        self.negative_samples = negative_samples
+
+    def forward(self, sentence_features: Iterable[Dict[str, torch.Tensor]], labels: torch.Tensor):
+        embeddings = [self.model(sentence_feature)['sentence_embedding'] for sentence_feature in sentence_features]
+        return self.multiple_negatives_ranking_loss(embeddings)
+
+    def multiple_negatives_ranking_loss(self, embeddings: List[torch.Tensor]):
+        positive_loss = torch.mean(torch.sum(embeddings[0] * embeddings[1], dim=-1))
+
+        negative_loss = torch.sum(embeddings[0] * embeddings[1], dim=-1)
+        for i in range(2, self.negative_samples + 2):
+            negative_loss = torch.cat((negative_loss, torch.sum(embeddings[0] * embeddings[i], dim=-1)), dim=-1)
+        negative_loss = torch.mean(torch.logsumexp(negative_loss, dim=-1))
+
+        return -positive_loss + negative_loss
+
+
+class MultipleNegativesRankingLoss(nn.Module):
+    def __init__(self, model: SentenceTransformer):
+        super(MultipleNegativesRankingLoss, self).__init__()
+        self.model = model
+
+    def forward(self, sentence_features: Iterable[Dict[str, torch.Tensor]], labels: torch.Tensor):
+        reps = [self.model(sentence_feature)['sentence_embedding'] for sentence_feature in sentence_features]
+
+        reps_a, reps_b = reps
+        return self.multiple_negatives_ranking_loss(reps_a, reps_b)
+
+    def multiple_negatives_ranking_loss(self, embeddings_a: torch.Tensor, embeddings_b: torch.Tensor):
+        scores = torch.matmul(embeddings_a, embeddings_b.t())
+        diagonal_mean = torch.mean(torch.diag(scores))
+        mean_log_row_sum_exp = torch.mean(torch.logsumexp(scores, dim=1))
+        return -diagonal_mean + mean_log_row_sum_exp
+
+
+class Pooling(nn.Module):
+    """Performs pooling (max or mean) on the token embeddings.
+
+    Using pooling, it generates from a variable sized sentence a fixed sized sentence embedding. This layer also allows to use the CLS token if it is returned by the underlying word embedding model.
+    You can concatenate multiple poolings together.
+    """
+    def __init__(self,
+                 word_embedding_dimension: int,
+                 pooling_mode_cls_token: bool = False,
+                 pooling_mode_max_tokens: bool = False,
+                 pooling_mode_mean_tokens: bool = True,
+                 pooling_mode_mean_sqrt_len_tokens: bool = False,
+                 ):
+        super(Pooling, self).__init__()
+
+        self.config_keys = ['word_embedding_dimension',  'pooling_mode_cls_token', 'pooling_mode_mean_tokens', 'pooling_mode_max_tokens', 'pooling_mode_mean_sqrt_len_tokens']
+
+        self.word_embedding_dimension = word_embedding_dimension
+        self.pooling_mode_cls_token = pooling_mode_cls_token
+        self.pooling_mode_mean_tokens = pooling_mode_mean_tokens
+        self.pooling_mode_max_tokens = pooling_mode_max_tokens
+        self.pooling_mode_mean_sqrt_len_tokens = pooling_mode_mean_sqrt_len_tokens
+
+        pooling_mode_multiplier = sum([pooling_mode_cls_token, pooling_mode_max_tokens, pooling_mode_mean_tokens, pooling_mode_mean_sqrt_len_tokens])
+        self.pooling_output_dimension = (pooling_mode_multiplier * word_embedding_dimension)
+
+    def forward(self, features: Dict[str, torch.Tensor]):
+        token_embeddings = features['token_embeddings']
+        cls_token = features['cls_token_embeddings']
+        attention_mask = features['attention_mask']
+
+        ## Pooling strategy
+        output_vectors = []
+        if self.pooling_mode_cls_token:
+            output_vectors.append(cls_token)
+        if self.pooling_mode_max_tokens:
+            input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+            token_embeddings[input_mask_expanded == 0] = -1e9  # Set padding tokens to large negative value
+            max_over_time = torch.max(token_embeddings, 1)[0]
+            output_vectors.append(max_over_time)
+        if self.pooling_mode_mean_tokens or self.pooling_mode_mean_sqrt_len_tokens:
+            input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+            sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
+
+            #If tokens are weighted (by WordWeights layer), feature 'token_weights_sum' will be present
+            if 'token_weights_sum' in features:
+                sum_mask = features['token_weights_sum'].unsqueeze(-1).expand(sum_embeddings.size())
+            else:
+                sum_mask = input_mask_expanded.sum(1)
+
+            sum_mask = torch.clamp(sum_mask, min=1e-9)
+
+            if self.pooling_mode_mean_tokens:
+                output_vectors.append(sum_embeddings / sum_mask)
+            if self.pooling_mode_mean_sqrt_len_tokens:
+                output_vectors.append(sum_embeddings / torch.sqrt(sum_mask))
+
+        output_vector = torch.cat(output_vectors, 1)
+        features.update({'sentence_embedding': output_vector})
+        return features
+
+    def get_sentence_embedding_dimension(self):
+        return self.pooling_output_dimension
+
+    def get_config_dict(self):
+        return {key: self.__dict__[key] for key in self.config_keys}
+
+    def save(self, output_path):
+        with open(os.path.join(output_path, 'config.json'), 'w') as fOut:
+            json.dump(self.get_config_dict(), fOut, indent=2)
+
+    @staticmethod
+    def load(input_path):
+        with open(os.path.join(input_path, 'config.json')) as fIn:
+            config = json.load(fIn)
+
+        return Pooling(**config)
+
+
+class RankingEvaluator:
+    def __init__(self, dataloader: DataLoader, random_paragraphs: DataLoader = None,
+                 name: str = '', show_progress_bar: bool = None):
+        """
+        Constructs an evaluator based for the dataset
+
+        The labels need to indicate the similarity between the sentences.
+
+        :param dataloader:
+            the data for the evaluation
+        :param main_similarity:
+            the similarity metric that will be used for the returned score
+        """
+        self.dataloader = dataloader
+        self.name = name
+        if name:
+            name = "_" + name
+
+        if show_progress_bar is None:
+            show_progress_bar = (
+                    logging.getLogger().getEffectiveLevel() == logging.INFO or logging.getLogger().getEffectiveLevel() == logging.DEBUG)
+        self.show_progress_bar = show_progress_bar
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.random_paragraphs = random_paragraphs
+
+    def __call__(self, model: 'SequentialSentenceEmbedder', output_path: str = None, epoch: int = -1,
+                 steps: int = -1) -> float:
+        model.eval()
+        embeddings1 = []
+        embeddings2 = []
+
+        if epoch != -1:
+            if steps == -1:
+                out_txt = f" after epoch {epoch}:"
+            else:
+                out_txt = f" in epoch {epoch} after {steps} steps:"
+        else:
+            out_txt = ":"
+
+        logging.info("Evaluation" + out_txt)
+
+        logging.info('Calculating sentence embeddings...')
+
+        self.dataloader.collate_fn = model.smart_batching_collate
+        iterator = self.dataloader
+        evidence_features = []
+        for step, batch in enumerate(tqdm(iterator, desc='batch')):
+            features, label_ids = batch_to_device(batch, self.device)
+            with torch.no_grad():
+                emb1, emb2 = [model(sent_features)['sentence_embedding'].to("cpu").numpy() for sent_features in
+                              features]
+            embeddings1.extend(emb1)
+            embeddings2.extend(emb2)
+
+            input_ids_2 = features[1]['input_ids'].to("cpu").numpy()
+            for f in input_ids_2:
+                evidence_features.append(hashlib.sha1(str(f).encode('utf-8')).hexdigest())
+
+        total_examples = len(embeddings1)
+
+        logging.info('Building index...')
+
+        vindex = VectorIndex(len(embeddings1[0]))
+        for v in embeddings2:
+            vindex.add(v)
+
+        if self.random_paragraphs is not None:
+            logging.info('Calculating random wikipedia paragraph embeddings...')
+
+            self.random_paragraphs.collate_fn = model.smart_batching_collate
+            iterator = self.random_paragraphs
+            for step, batch in enumerate(iterator):
+                if step % 10 == 0:
+                    logging.info('Batch {}/{}'.format(step, len(iterator)))
+                features, label_ids = batch_to_device(batch, self.device)
+                with torch.no_grad():
+                    embeddings = model(features[0])['sentence_embedding'].to("cpu").numpy()
+                for emb in embeddings:
+                    vindex.add(emb)
+
+        vindex.build()
+
+        logging.info('Ranking evaluation...')
+
+        mrr = 1e-8
+        recall = {1: 0, 5: 0, 10: 0, 20: 0, 100: 0}
+
+        all_results, _ = vindex.search(embeddings1, k=100, probes=1024)
+
+        for i in range(total_examples):
+            results = all_results[i]
+            rank = 1
+            found = False
+            for r in results:
+                if r < len(evidence_features) and evidence_features[r] == evidence_features[i]:
+                    mrr += 1 / rank
+                    found = True
+                    break
+                rank += 1
+
+            for topk, count in recall.items():
+                if rank <= topk and found:
+                    recall[topk] += 1
+        mrr /= total_examples
+        for topk, count in recall.items():
+            recall[topk] /= total_examples
+            logging.info('recall@{} : {}'.format(topk, recall[topk]))
+
+        logging.info('mrr@100 : {}'.format(mrr))
+
+        if output_path is not None:
+            f = open(output_path + '/stats.csv', 'a+')
+            f.write('{},{},{},{},{},{}\n'.format(recall[1], recall[5], recall[10], recall[20], recall[100], mrr))
+            f.close()
+
+        return mrr
+
+
+if __name__ == "__main__":
+    logging.basicConfig(format='%(asctime)s - %(message)s',
+                        datefmt='%Y-%m-%d %H:%M:%S',
+                        level=logging.CRITICAL,
+                        handlers=[LoggingHandler()])
+
+    # Read the dataset
+    model_name = 'distilroberta-base'
+    batch_size = 384
+    model_save_path = 'weights/encoder/qrbert-multitask-distil'
+    num_epochs = 1000
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Use BERT for mapping tokens to embeddings
+    word_embedding_model = BERT(model_name, max_seq_length=256, do_lower_case=False)
+
+    # Apply mean pooling to get one fixed sized sentence vector
+    pooling_model = Pooling(word_embedding_model.get_word_embedding_dimension(),
+                            pooling_mode_mean_tokens=True,
+                            pooling_mode_cls_token=False,
+                            pooling_mode_max_tokens=False)
+
+    if torch.cuda.device_count() > 1:
+        print("Let's use", torch.cuda.device_count(), "GPUs!")
+
+    word_embedding_model = DataParallel(word_embedding_model)
+    word_embedding_model.to(device)
+
+    model = SentenceTransformer(modules=[word_embedding_model, pooling_model], parallel=True)
+
+    #model = SentenceTransformer('models/weights/encoder/sbert-nli-fnli-qar2-768', parallel=True)
+
+    logging.info("Reading dev dataset")
+    dev_data = SentencesDataset(get_retrieval_examples(filename='../data/retrieval/dev.jsonl',
+                                                       no_statements=False,
+                                                       ), model=model)
+    dev_dataloader = DataLoader(dev_data, shuffle=True, batch_size=1024)
+    evaluator = RankingEvaluator(dev_dataloader)
+    #model.best_score = model.evaluate(evaluator)
+
+    train_data = SentencesDataset(get_retrieval_examples(filename='../data/retrieval/train.jsonl',
+                                                         no_statements=False),
+                                  model=model)
+    train_dataloader = DataLoader(train_data, shuffle=True, batch_size=batch_size)
+    train_loss = MultipleNegativesRankingLoss(model=model)
+
+    for i in range(num_epochs):
+        logging.info("Epoch {}".format(i))
+        # Train the model
+        model.fit(train_objectives=[(train_dataloader, train_loss)],
+                  evaluator=evaluator,
+                  epochs=1,
+                  acc_steps=1,
+                  evaluation_steps=1000,
+                  warmup_steps=1000,
+                  output_path=model_save_path,
+                  optimizer_params={'lr': 1e-6, 'eps': 1e-6, 'correct_bias': False}
+                  )
+        torch.cuda.empty_cache()

models/tokenization.py

@@ -0,0 +1,56 @@
+import numpy as np
+
+
+def get_segments(sentence):
+    sentence_segments = []
+    temp = []
+    i = 0
+    for token in sentence.split(" "):
+        temp.append(i)
+        if token == "[SEP]":
+            i += 1
+    sentence_segments.append(temp)
+    return sentence_segments
+
+
+def tokenize(text, max_length, tokenizer, second_text=None):
+    if second_text is None:
+        sentence = "[CLS] " + " ".join(tokenizer.tokenize(text.replace('[^\w\s]+|\n', ''))[:max_length-2]) + " [SEP]"
+    else:
+        text = tokenizer.tokenize(text.replace('[^\w\s]+|\n', ''))
+        second_text = tokenizer.tokenize(second_text.replace('[^\w\s]+|\n', ''))
+        while len(text) + len(second_text) > max_length - 3:
+            if len(text) > len(second_text):
+                text.pop()
+            else:
+                second_text.pop()
+        sentence = "[CLS] " + " ".join(text) + " [SEP] " + " ".join(second_text) + " [SEP]"
+
+    # generate masks
+    # bert requires a mask for the words which are padded.
+    # Say for example, maxlen is 100, sentence size is 90. then, [1]*90 + [0]*[100-90]
+    sentence_mask = [1] * len(sentence.split(" ")) + [0] * (max_length - len(sentence.split(" ")))
+
+    # generate input ids
+    # if less than max length provided then the words are padded
+    if len(sentence.split(" ")) != max_length:
+        sentence_padded = sentence + " [PAD]" * (max_length - len(sentence.split(" ")))
+    else:
+        sentence_padded = sentence
+
+    sentence_converted = tokenizer.convert_tokens_to_ids(sentence_padded.split(" "))
+
+    # generate segments
+    # for each separation [SEP], a new segment is converted
+    sentence_segment = get_segments(sentence_padded)
+
+    # convert list into tensor integer arrays and return it
+    # return sentences_converted,sentences_segment, sentences_mask
+    """
+    return [tf.cast(sentence_converted, tf.int32),
+            tf.cast(sentence_segment, tf.int32),
+            tf.cast(sentence_mask, tf.int32)]
+    """
+    return [np.asarray(sentence_converted, dtype=np.int32).squeeze(),
+            np.asarray(sentence_segment, dtype=np.int32).squeeze(),
+            np.asarray(sentence_mask, dtype=np.int32).squeeze()]

models/vector_index.py

@@ -0,0 +1,102 @@
+import logging
+import math
+import pickle
+import h5py
+import faiss
+
+import numpy as np
+import gpustat
+
+import pdb
+
+# see http://ulrichpaquet.com/Papers/SpeedUp.pdf theorem 5
+
+def get_phi(xb):
+    return (xb ** 2).sum(1).max()
+
+
+def augment_xb(xb, phi=None):
+    norms = (xb ** 2).sum(1)
+    if phi is None:
+        phi = norms.max()
+    extracol = np.sqrt(phi - norms)
+    return np.hstack((xb, extracol.reshape(-1, 1)))
+
+
+def augment_xq(xq):
+    extracol = np.zeros(len(xq), dtype='float32')
+    return np.hstack((xq, extracol.reshape(-1, 1)))
+
+
+class VectorIndex:
+    def __init__(self, d):
+        self.d = d
+        self.vectors = []
+        self.index = None
+
+    def add(self, v):
+        self.vectors.append(v)
+
+    def add_vectors(self, vs):
+        logging.info('Adding vectors to index...')
+        self.index.add(vs)
+
+
+    def build(self, use_gpu=False):
+        self.vectors = np.array(self.vectors) # OOM at this step if building too many vectors
+
+        faiss.normalize_L2(self.vectors)
+
+        #self.vectors = augment_xq(self.vectors)
+
+        logging.info('Indexing {} vectors'.format(self.vectors.shape[0]))
+
+        if self.vectors.shape[0] > 50000:
+            num_centroids = 8 * int(math.sqrt(math.pow(2, int(math.log(self.vectors.shape[0], 2)))))
+
+            logging.info('Using {} centroids'.format(num_centroids))
+
+            self.index = faiss.index_factory(self.d, "IVF{}_HNSW32,Flat".format(num_centroids))
+
+            ngpu = faiss.get_num_gpus()
+            if ngpu > 0 and use_gpu:
+                logging.info('Using {} GPUs'.format(ngpu))
+
+                index_ivf = faiss.extract_index_ivf(self.index)
+                gpustat.print_gpustat()
+                clustering_index = faiss.index_cpu_to_all_gpus(faiss.IndexFlatL2(self.d))
+                index_ivf.clustering_index = clustering_index
+
+            logging.info('Training index...')
+
+            self.index.train(self.vectors)
+        else:
+            self.index = faiss.IndexFlatL2(self.d)
+            if faiss.get_num_gpus() > 0 and use_gpu:
+                gpustat.print_gpustat()
+                self.index = faiss.index_cpu_to_all_gpus(self.index)
+
+    def load(self, path):
+        self.index = faiss.read_index(path)
+
+    def save(self, path):
+        gpustat.print_gpustat()
+        faiss.write_index(faiss.index_gpu_to_cpu(self.index), path)
+
+    def save_vectors(self, path):
+        #pickle.dump(self.vectors, open(path, 'wb'), protocol=4)
+        f = h5py.File(path, 'w')
+        dset = f.create_dataset('data', data=self.vectors)
+        f.close()
+
+    def search(self, vectors, k=1, probes=1):
+        if not isinstance(vectors, np.ndarray):
+            vectors = np.array(vectors)
+        #faiss.normalize_L2(vectors)
+        try:
+            self.index.nprobe = probes
+        except:
+            pass
+        distances, ids = self.index.search(vectors, k)
+        similarities = [(2-d)/2 for d in distances]
+        return ids, similarities

requirements.txt

@@ -0,0 +1,30 @@
+tqdm==4.41.0
+transformers==3.0.2
+faiss_gpu==1.6.3
+syllables==0.1.0
+scipy==1.3.1
+nltk==3.2.5
+lxml==4.6.3
+requests==2.22.0
+Flask_Cors==3.0.8
+gpustat==0.6.0
+h5py==2.10.0
+packaging==20.3
+pandas==1.0.3
+spacy==3.0.3
+sentence-transformers==0.3.0
+numpy==1.17.4
+matplotlib==3.3.1
+filelock==3.0.12
+boilerpipe3==1.1
+Unidecode==1.1.1
+Flask==2.0.0
+apex
+beautifulsoup4==4.9.3
+python_dateutil==2.8.1
+scikit_learn==0.24.1
+tantivy==0.13.2
+huggingface_hub==0.16.4
+torch==1.6.0+cu101
+torchvision==0.7.0+cu101
+gradio==3.34.0

utils.py

@@ -0,0 +1,71 @@
+from html.parser import HTMLParser
+
+
+class WebParser(HTMLParser):
+    """
+    A class for converting the tagged html to formats that can be used by a ML model
+    """
+    def __init__(self):
+        super().__init__()
+        self.block_tags = {
+            'div', 'p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6'
+        }
+        self.inline_tags = {
+            '', 'a', 'b', 'main', 'span', 'em', 'strong', 'br'
+        }
+        self.allowed_tags = {'div', 'p', '', 'a', 'b', 'main', 'span', 'em', 'strong', 'br'}
+
+        self.opened_tags = []
+        self.block_content = ''
+        self.blocks = []
+
+    def get_last_opened_tag(self):
+        """
+        Gets the last visited tag
+        :return:
+        """
+        if len(self.opened_tags) > 0:
+            return self.opened_tags[len(self.opened_tags) - 1]
+        return ''
+
+    def error(self, message):
+        pass
+
+    def handle_starttag(self, tag, attrs):
+        """
+        Handles the start tag of an HTML node in the tree
+        :param tag: the HTML tag
+        :param attrs: the tag attributes
+        :return:
+        """
+        self.opened_tags.append(tag)
+
+    def handle_endtag(self, tag):
+        """
+        Handles the end tag of an HTML node in the tree
+        :param tag: the HTML tag
+        :return:
+        """
+        if tag in self.block_tags:
+            self.block_content = self.block_content.strip()
+            if len(self.block_content) > 0:
+                #if not self.block_content.endswith('.'): self.block_content += '.'
+                self.blocks.append(self.block_content)
+            self.block_content = ''
+        if len(self.opened_tags) > 0:
+            self.opened_tags.pop()
+
+    def handle_data(self, data):
+        """
+        Handles a text HTML node in the tree
+        :param data: the text node
+        :return:
+        """
+        last_opened_tag = self.get_last_opened_tag()
+        if last_opened_tag in self.allowed_tags:
+            data = data.replace('  ', ' ').strip()
+            if data != '':
+                self.block_content += data + ' '
+
+    def get_blocks(self):
+        return self.blocks

web_search.py

@@ -0,0 +1,140 @@
+import re
+
+from pprint import pprint
+from urllib.request import Request, urlopen
+from html.parser import HTMLParser
+from bs4 import BeautifulSoup
+from urllib.parse import quote_plus
+
+
+def bing_search(query: str, pages_number=1) -> list:
+    """
+    Gets web results from Bing
+    :param query: query to search
+    :param pages_number: number of search pages to scrape
+    :return: a list of links in ranked order
+    """
+    urls = []
+    for page in range(pages_number):
+        first = page * 10 + 1
+        address = "https://www.bing.com/search?q=" + quote_plus(query) + '&first=' + str(first)
+        data = get_html(address)
+        soup = BeautifulSoup(data, 'lxml')
+        links = soup.findAll('li', {'class': 'b_algo'})
+        urls.extend([link.find('h2').find('a')['href'] for link in links])
+
+    return urls
+
+
+def duckduckgo_search(query: str, pages=1):
+    urls = []
+    start_index = 0
+    for page in range(pages):
+        address = "https://duckduckgo.com/html/?kl=en-us&q={}&s={}".format(quote_plus(query), start_index)
+        data = get_html(address)
+        soup = BeautifulSoup(data, 'lxml')
+        links = soup.findAll('a', {'class': 'result__a'})
+        urls.extend([link['href'] for link in links])
+        start_index = len(urls)
+
+    return urls
+
+
+def news_search(query: str, pages=1):
+    urls = []
+    for page in range(pages):
+        api_url = f'https://newslookup.com/results?l=2&q={quote_plus(query)}&dp=&mt=-1&mkt=0&mtx=0&mktx=0&s=&groupby=no&cat=-1&from=&fmt=&tp=720&ps=50&ovs=&page={page}'
+        data = get_html(api_url)
+        soup = BeautifulSoup(data, 'lxml')
+        links = soup.findAll('a', {'class': 'title'})
+        urls.extend([link['href'] for link in links])
+    return urls
+
+
+def get_html(url: str) -> str:
+    """
+    Downloads the html source code of a webpage
+    :param url:
+    :return: html source code
+    """
+    try:
+        custom_user_agent = "Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:47.0) Gecko/20100101 Firefox/47.0"
+        req = Request(url, headers={"User-Agent": custom_user_agent})
+        page = urlopen(req, timeout=3)
+        return str(page.read())
+    except:
+        return ''
+
+
+class WebParser(HTMLParser):
+    """
+    A class for converting the tagged html to formats that can be used by a ML model
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.block_tags = {
+            'div', 'p'
+        }
+        self.inline_tags = {
+            '', 'a', 'b', 'tr', 'main', 'span', 'time', 'td',
+            'sup', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'em', 'strong', 'br'
+        }
+        self.allowed_tags = self.block_tags.union(self.inline_tags)
+        self.opened_tags = []
+        self.block_content = ''
+        self.blocks = []
+
+    def get_last_opened_tag(self):
+        """
+        Gets the last visited tag
+        :return:
+        """
+        if len(self.opened_tags) > 0:
+            return self.opened_tags[len(self.opened_tags) - 1]
+        return ''
+
+    def error(self, message):
+        pass
+
+    def handle_starttag(self, tag, attrs):
+        """
+        Handles the start tag of an HTML node in the tree
+        :param tag: the HTML tag
+        :param attrs: the tag attributes
+        :return:
+        """
+        self.opened_tags.append(tag)
+        if tag in self.block_tags:
+            self.block_content = self.block_content.strip()
+            if len(self.block_content) > 0:
+                if not self.block_content.endswith('.'):
+                    self.block_content += '.'
+                self.block_content = self.block_content.replace('\\n', ' ').replace('\\r', ' ')
+                self.block_content = re.sub("\s\s+", " ", self.block_content)
+                self.blocks.append(self.block_content)
+            self.block_content = ''
+
+    def handle_endtag(self, tag):
+        """
+        Handles the end tag of an HTML node in the tree
+        :param tag: the HTML tag
+        :return:
+        """
+        if len(self.opened_tags) > 0:
+            self.opened_tags.pop()
+
+    def handle_data(self, data):
+        """
+        Handles a text HTML node in the tree
+        :param data: the text node
+        :return:
+        """
+        last_opened_tag = self.get_last_opened_tag()
+        if last_opened_tag in self.allowed_tags:
+            data = data.replace('  ', ' ').strip()
+            if data != '':
+                self.block_content += data + ' '
+
+    def get_text(self):
+        return "\n\n".join(self.blocks)