remove all files but the indices

.gitattributes

@@ -1,39 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-*.bib filter=lfs diff=lfs merge=lfs -text
-collection.json filter=lfs diff=lfs merge=lfs -text
-dataset.json filter=lfs diff=lfs merge=lfs -text
-index/metadata.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,4 +0,0 @@
-__pycache__
-experiments
-.DS_Store
-app

Dockerfile

@@ -1,18 +0,0 @@
-FROM python:3.10
-
-WORKDIR /app
-
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-COPY . .
-
-# Copy ColBERT files that aren't downloaded properly
-COPY ./src/extras/segmented_maxsim.cpp /usr/local/lib/python3.10/site-packages/colbert/modeling/segmented_maxsim.cpp
-COPY ./src/extras/decompress_residuals.cpp /usr/local/lib/python3.10/site-packages/colbert/search/decompress_residuals.cpp
-COPY ./src/extras/filter_pids.cpp /usr/local/lib/python3.10/site-packages/colbert/search/filter_pids.cpp
-COPY ./src/extras/segmented_lookup.cpp /usr/local/lib/python3.10/site-packages/colbert/search/segmented_lookup.cpp
-
-# CMD ["sh", "-c", "sleep infinity"]
-CMD ["python", "src/server.py"]
-# CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:8893", "src/server:app"]

README.md

@@ -2,110 +2,6 @@
 license: apache-2.0
 ---
 
-Use ColBERT as a search engine for the [ACL Anthology](https://aclanthology.org/). (Parse any bibtex, and store in a MySQL service)
+Use ColBERT as a search engine for the [ACL Anthology](https://aclanthology.org/). (Parse any bibtex, and store in a MySQL service).
 
-# Setup
-
-## Setup ColBERT
-```sh
-git clone https://huggingface.co/davidheineman/colbert-acl
-
-# install dependencies
-# torch==1.13.1 required (conda install -y -n [env] python=3.10)
-pip install -r requirements.txt
-brew install mysql
-```
-
-### (Optional) Parse & Index the Anthology
-
-Feel free to skip, since the parsed/indexed anthology is contained in this repo. 
-
-```sh
-# get up-to-date abstracts in bibtex
-curl -O https://aclanthology.org/anthology+abstracts.bib.gz
-gunzip anthology+abstracts.bib.gz
-mv anthology+abstracts.bib anthology.bib
-
-# parse .bib -> .json
-python parse.py
-
-# index with ColBERT 
-# (note sometimes there is a silent failure if the CPP extensions do not exist)
-python index.py
-```
-
-### Search with ColBERT
-
-```sh
-# start flask server
-python server.py
-
-# or start a production API endpoint
-gunicorn -w 4 -b 0.0.0.0:8893 server:app
-```
-
-Then, to test, visit:
-```
-http://localhost:8893/api/search?query=Information retrevial with BERT
-```
-or for an interface:
-```
-http://localhost:8893
-```
-
-### Deploy as a Docker App
-```sh
-docker-compose build --no-cache
-docker-compose up --build
-```
-
-## Example notebooks
-
-To see an example of search, visit:
-[colab.research.google.com/drive/1-b90_8YSAK17KQ6C7nqKRYbCWEXQ9FGs](https://colab.research.google.com/drive/1-b90_8YSAK17KQ6C7nqKRYbCWEXQ9FGs?usp=sharing)
-
-<!-- ## Notes
-- See: 
-    - https://github.com/stanford-futuredata/ColBERT/blob/main/colbert/index_updater.py
-    - https://github.com/stanford-futuredata/ColBERT/issues/111
-
-- TODO:
-    - On UI
-        - Colors: make the colors resemble the ACL page much closer
-            - There's still a bunch of blue from the bootstrap themeing
-        - Smaller line spacing for abstract text
-        - Add "PDF" button
-        - Justify the result metadata (Year, venue, etc.) so the content all starts at the same vertical position
-        - Add a "Expand" button at the end of the abstract
-        - Make the results scrollable, without scrolling the rest of the page
-        - Put two sliders on the year range (and make the years selectable, with the years at both ends of the bar)
-        - If the user selects certain venues, remember these venues
-        - Add a dropdown under the "Workshop" box to select specific workshops
-
-    - Include the title in the indexing
-    
-    - https://docs.docker.com/language/python/configure-ci-cd/
-
-    - Have articles before 2020
-
-    - Maybe make the UI more compressed like this: https://aclanthology.org/events/eacl-2024/#2024eacl-long
-
-    - Put query in URL (?q=XXX)
-
-    - Move code to github and index to hf, then use this to download the index:
-        from huggingface_hub import snapshot_download
-
-        # Download indexed repo at: https://huggingface.co/davidheineman/colbert-acl
-        !mkdir "acl"
-        index_name = snapshot_download(repo_id="davidheineman/colbert-acl", local_dir="acl")
-    - Make indexing much easier 
-        (currently, the setup involves manually copying the CPP files becuase there is a silent failure, this also should be possible to do on Google Collab, or even MPS)
-        - Make index save in parent folder
-        - Fix "sanity check" in index.py
-    - Profile bibtexparser.load(f) (why so slow)
-    - Ship as a containerized service
-    - Scrape: 
-        - https://proceedings.neurips.cc/
-        - https://dblp.uni-trier.de/db/conf/iclr/index.html
-        - openreview
- -->
+**This repo contains the ColBERT index and dataset.** To run the interface, see [github.com/davidheineman/acl-search](https://github.com/davidheineman/acl-search).

docker-compose.yml

@@ -1,33 +0,0 @@
-services:
-  mysql:
-    image: mysql:8.0
-    container_name: mysql_db
-    environment:
-      MYSQL_ROOT_PASSWORD: 
-      MYSQL_ALLOW_EMPTY_PASSWORD: true
-      MYSQL_DATABASE: anthology
-      MYSQL_USER: myuser
-      MYSQL_PASSWORD: mysecret
-    volumes:
-      - mysql_data:/var/lib/mysql
-    networks:
-      - mysql_network
-
-  python:
-    build:
-      context: .
-      dockerfile: Dockerfile
-    container_name: python_app
-    ports:
-      - "8893:8893" # host:local
-    depends_on:
-      - mysql
-    networks:
-      - mysql_network
-
-networks:
-  mysql_network:
-    driver: bridge
-
-volumes:
-  mysql_data:

requirements.txt

@@ -1,7 +0,0 @@
-torch==1.13.1
-colbert-ir[torch,faiss-cpu] # faiss-gpu
-faiss-cpu # shouldn't have to include
-bibtexparser
-mysql-connector-python
-flask
-gunicorn

src/constants.py

@@ -1,15 +0,0 @@
-import os
-from typing import Literal
-
-INDEX_NAME = os.getenv("INDEX_NAME", 'index')
-INDEX_ROOT = os.getenv("INDEX_ROOT", os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-INDEX_PATH = os.path.join(INDEX_ROOT, INDEX_NAME)
-ANTHOLOGY_PATH = os.path.join(INDEX_ROOT, 'anthology.bib')
-DATASET_PATH = os.path.join(INDEX_ROOT, 'dataset.json')
-
-DB_NAME = 'anthology'
-DB_HOSTNAME = 'mysql_db' # localhost
-DB_PORT = 3306 # None
-
-VENUES = Literal['workshop', 'journal', 'short', 'demo', 'tutorial', 'industry', 'findings', 'main']

src/db.py

@@ -1,166 +0,0 @@
-import json
-from typing import List, Optional, Union
-
-import mysql.connector
-from constants import DATASET_PATH, DB_HOSTNAME, DB_NAME, DB_PORT, VENUES
-
-PAPER_QUERY = """
-SELECT * 
-FROM paper 
-WHERE pid IN ({query_arg_str}){constraints_str};
-"""
-
-def read_dataset():
-    print("Reading dataset...")
-    with open(DATASET_PATH, 'r', encoding='utf-8') as f:
-        dataset = json.loads(f.read())
-    dataset = [d for d in dataset if 'abstract' in d.keys()]
-    return dataset
-
-
-def create_database():
-    db = mysql.connector.connect(
-        host = DB_HOSTNAME,
-        user = "root",
-        password = "",
-        port = DB_PORT
-    )
-    cursor = db.cursor()
-
-    cursor.execute("SHOW DATABASES")
-    db_exists = False
-    for x in cursor:
-        db_name = x[0]
-        if db_name == DB_NAME:
-            db_exists = True
-
-    # Create database
-    if not db_exists:
-        print("Creating new database...")
-        cursor.execute(f'CREATE DATABASE {DB_NAME}')
-    cursor.execute(f'USE {DB_NAME}')
-
-    # Create table
-    print('Creating new table...')
-    cursor.execute(f'DROP TABLE IF EXISTS paper')
-    cursor.execute("""
-    CREATE TABLE paper (
-        pid INT PRIMARY KEY, 
-        title VARCHAR(1024), 
-        author VARCHAR(2170), 
-        year INT, 
-        abstract TEXT(12800), 
-        url VARCHAR(150), 
-        type VARCHAR(100), 
-        venue VARCHAR(500), 
-        venue_type VARCHAR(150),
-        is_findings TINYINT(1) NOT NULL DEFAULT 0
-    )
-    """)
-
-    acl_data = read_dataset()
-
-    vals = []
-    for pid, paper in enumerate(acl_data):
-        title    = paper.get('title', '')
-        author   = paper.get('author', '')
-        year     = paper.get('year', '')
-        abstract = paper.get('abstract', '')
-        url      = paper.get('url', '')
-        type     = paper.get('ENTRYTYPE', '')
-        venue    = paper.get('booktitle', '')
-        venue_type  = paper.get('venue_type', '')
-        is_findings = paper.get('is_findings', '0')
-
-        if not abstract: continue
-
-        vals += [(pid, title, author, year, abstract, url, type, venue, venue_type, is_findings)]
-
-    sql = """
-    INSERT INTO paper (
-        pid, title, author, year, abstract, url, type, venue, venue_type, is_findings
-    ) VALUES (
-        %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
-    )
-    """
-
-    print('Writing entries to table...')
-    cursor.executemany(sql, vals)
-    db.commit()
-
-
-def parse_results(results):
-    parsed_results = {}
-
-    for result in results:
-        pid, title, authors, year, abstract, url, type, venue, venue_type, is_findings = result
-
-        title    = title.replace("{", "").replace("}", "")
-        authors  = authors.replace("{", "").replace("}", "").replace('\\"', "")
-        abstract = abstract.replace("{", "").replace("}", "").replace("\\", "")
-
-        parsed_results[int(pid)] = {
-            'title': title, 
-            'authors': authors, 
-            'year': year, 
-            'abstract': abstract,
-            'url': url,
-            'type': type,
-            'venue': venue,
-            'venue_type': venue_type,
-            'is_findings': is_findings,
-        }
-
-    return parsed_results
-
-
-def query_paper_metadata(
-        pids: List[int], 
-        start_year: int = None, 
-        end_year: int = None, 
-        venue_type: Union[VENUES, List[VENUES]] = None, 
-        is_findings: Optional[bool] = None
-    ):
-    if not isinstance(venue_type, list): venue_type = [venue_type]
-    
-    db = mysql.connector.connect(
-        host = DB_HOSTNAME,
-        user = "root",
-        password = "",
-        database = DB_NAME,
-        port = DB_PORT
-    )
-
-    cursor = db.cursor()
-
-    pids_str = ', '.join(['%s'] * len(pids))
-
-    constraints_str = ""
-    if start_year: constraints_str += f" AND year >= {start_year}"
-    if end_year: constraints_str += f" AND year <= {end_year}"
-    if is_findings: constraints_str += f" AND is_findings = {is_findings}"
-    if venue_type: 
-        venue_str = ','.join([f'"{venue}"' for venue in venue_type])
-        constraints_str += f" AND venue_type IN ({venue_str})"
-
-    query = PAPER_QUERY.format(
-        query_arg_str=pids_str, 
-        constraints_str=constraints_str
-    )
-
-    # print(PAPER_QUERY.format(query_arg_str=', '.join([str(p) for p in pids]), year=year))
-    
-    cursor.execute(query, pids)
-    results = cursor.fetchall()
-    
-    if len(results) == 0: return []
-
-    parsed_results = parse_results(results)
-
-    # Restore original ordering of PIDs from ColBERT
-    results = [parsed_results[pid] for pid in pids if pid in parsed_results.keys()]
-
-    return results
-
-
-if __name__ == '__main__': create_database()

src/extras/decompress_residuals.cpp

@@ -1,160 +0,0 @@
-#include <pthread.h>
-#include <torch/extension.h>
-
-typedef struct decompress_args {
-    int tid;
-    int nthreads;
-
-    int npids;
-    int dim;
-    int packed_dim;
-    int npacked_vals_per_byte;
-
-    int* pids;
-    int64_t* lengths;
-    int64_t* offsets;
-    float* bucket_weights;
-    uint8_t* reversed_bit_map;
-    uint8_t* bucket_weight_combinations;
-    uint8_t* binary_residuals;
-    int* codes;
-    float* centroids;
-    int64_t* cumulative_lengths;
-
-    float* output;
-} decompress_args_t;
-
-void* decompress(void* args) {
-    decompress_args_t* decompress_args = (decompress_args_t*)args;
-
-    int npids_per_thread = (int)std::ceil(((float)decompress_args->npids) /
-                                          decompress_args->nthreads);
-    int start = decompress_args->tid * npids_per_thread;
-    int end = std::min((decompress_args->tid + 1) * npids_per_thread,
-                       decompress_args->npids);
-
-    // Iterate over all documents
-    for (int i = start; i < end; i++) {
-        int pid = decompress_args->pids[i];
-
-        // Offset into packed list of token vectors for the given document
-        int64_t offset = decompress_args->offsets[pid];
-
-        // For each document, iterate over all token vectors
-        for (int j = 0; j < decompress_args->lengths[pid]; j++) {
-            const int code = decompress_args->codes[offset + j];
-
-            // For each token vector, iterate over the packed (8-bit) residual
-            // values
-            for (int k = 0; k < decompress_args->packed_dim; k++) {
-                uint8_t x =
-                    decompress_args->binary_residuals
-                        [(offset + j) * decompress_args->packed_dim + k];
-                x = decompress_args->reversed_bit_map[x];
-
-                // For each packed residual value, iterate over the bucket
-                // weight indices. If we use n-bit compression, that means there
-                // will be (8 / n) indices per packed value.
-                for (int l = 0; l < decompress_args->npacked_vals_per_byte;
-                     l++) {
-                    const int output_dim_idx =
-                        k * decompress_args->npacked_vals_per_byte + l;
-                    const int bucket_weight_idx =
-                        decompress_args->bucket_weight_combinations
-                            [x * decompress_args->npacked_vals_per_byte + l];
-                    decompress_args
-                        ->output[(decompress_args->cumulative_lengths[i] + j) *
-                                     decompress_args->dim +
-                                 output_dim_idx] =
-                        decompress_args->bucket_weights[bucket_weight_idx] +
-                        decompress_args->centroids[code * decompress_args->dim +
-                                                   output_dim_idx];
-                }
-            }
-        }
-    }
-
-    return NULL;
-}
-
-torch::Tensor decompress_residuals(
-    const torch::Tensor pids, const torch::Tensor lengths,
-    const torch::Tensor offsets, const torch::Tensor bucket_weights,
-    const torch::Tensor reversed_bit_map,
-    const torch::Tensor bucket_weight_combinations,
-    const torch::Tensor binary_residuals, const torch::Tensor codes,
-    const torch::Tensor centroids, const int dim, const int nbits) {
-    const int npacked_vals_per_byte = (8 / nbits);
-    const int packed_dim = (int)(dim / npacked_vals_per_byte);
-
-    int npids = pids.size(0);
-    int* pids_a = pids.data_ptr<int>();
-    int64_t* lengths_a = lengths.data_ptr<int64_t>();
-    int64_t* offsets_a = offsets.data_ptr<int64_t>();
-    float* bucket_weights_a = bucket_weights.data_ptr<float>();
-    uint8_t* reversed_bit_map_a = reversed_bit_map.data_ptr<uint8_t>();
-    uint8_t* bucket_weight_combinations_a =
-        bucket_weight_combinations.data_ptr<uint8_t>();
-    uint8_t* binary_residuals_a = binary_residuals.data_ptr<uint8_t>();
-    int* codes_a = codes.data_ptr<int>();
-    float* centroids_a = centroids.data_ptr<float>();
-
-    int64_t cumulative_lengths[npids + 1];
-    int noutputs = 0;
-    cumulative_lengths[0] = 0;
-    for (int i = 0; i < npids; i++) {
-        noutputs += lengths_a[pids_a[i]];
-        cumulative_lengths[i + 1] =
-            cumulative_lengths[i] + lengths_a[pids_a[i]];
-    }
-
-    auto options =
-        torch::TensorOptions().dtype(torch::kFloat32).requires_grad(false);
-    torch::Tensor output = torch::zeros({noutputs, dim}, options);
-    float* output_a = output.data_ptr<float>();
-
-    auto nthreads = at::get_num_threads();
-
-    pthread_t threads[nthreads];
-    decompress_args_t args[nthreads];
-
-    for (int i = 0; i < nthreads; i++) {
-        args[i].tid = i;
-        args[i].nthreads = nthreads;
-
-        args[i].npids = npids;
-        args[i].dim = dim;
-        args[i].packed_dim = packed_dim;
-        args[i].npacked_vals_per_byte = npacked_vals_per_byte;
-
-        args[i].pids = pids_a;
-        args[i].lengths = lengths_a;
-        args[i].offsets = offsets_a;
-        args[i].bucket_weights = bucket_weights_a;
-        args[i].reversed_bit_map = reversed_bit_map_a;
-        args[i].bucket_weight_combinations = bucket_weight_combinations_a;
-        args[i].binary_residuals = binary_residuals_a;
-        args[i].codes = codes_a;
-        args[i].centroids = centroids_a;
-        args[i].cumulative_lengths = cumulative_lengths;
-
-        args[i].output = output_a;
-
-        int rc = pthread_create(&threads[i], NULL, decompress, (void*)&args[i]);
-        if (rc) {
-            fprintf(stderr, "Unable to create thread %d: %d\n", i, rc);
-            std::exit(1);
-        }
-    }
-
-    for (int i = 0; i < nthreads; i++) {
-        pthread_join(threads[i], NULL);
-    }
-
-    return output;
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("decompress_residuals_cpp", &decompress_residuals,
-          "Decompress residuals", py::call_guard<py::gil_scoped_release>());
-}

src/extras/filter_pids.cpp

@@ -1,174 +0,0 @@
-#include <pthread.h>
-#include <torch/extension.h>
-
-#include <algorithm>
-#include <chrono>
-#include <numeric>
-#include <utility>
-
-typedef struct maxsim_args {
-    int tid;
-    int nthreads;
-
-    int ncentroids;
-    int nquery_vectors;
-    int npids;
-
-    int* pids;
-    float* centroid_scores;
-    int* codes;
-    int64_t* doclens;
-    int64_t* offsets;
-    bool* idx;
-
-    std::priority_queue<std::pair<float, int>> approx_scores;
-} maxsim_args_t;
-
-void* maxsim(void* args) {
-    maxsim_args_t* maxsim_args = (maxsim_args_t*)args;
-
-    float per_doc_approx_scores[maxsim_args->nquery_vectors];
-    for (int k = 0; k < maxsim_args->nquery_vectors; k++) {
-        per_doc_approx_scores[k] = -9999;
-    }
-
-    int ndocs_per_thread =
-        (int)std::ceil(((float)maxsim_args->npids) / maxsim_args->nthreads);
-    int start = maxsim_args->tid * ndocs_per_thread;
-    int end =
-        std::min((maxsim_args->tid + 1) * ndocs_per_thread, maxsim_args->npids);
-
-    std::unordered_set<int> seen_codes;
-
-    for (int i = start; i < end; i++) {
-        auto pid = maxsim_args->pids[i];
-        for (int j = 0; j < maxsim_args->doclens[pid]; j++) {
-            auto code = maxsim_args->codes[maxsim_args->offsets[pid] + j];
-            assert(code < maxsim_args->ncentroids);
-            if (maxsim_args->idx[code] &&
-                seen_codes.find(code) == seen_codes.end()) {
-                for (int k = 0; k < maxsim_args->nquery_vectors; k++) {
-                    per_doc_approx_scores[k] =
-                        std::max(per_doc_approx_scores[k],
-                                 maxsim_args->centroid_scores
-                                     [code * maxsim_args->nquery_vectors + k]);
-                }
-                seen_codes.insert(code);
-            }
-        }
-        float score = 0;
-        for (int k = 0; k < maxsim_args->nquery_vectors; k++) {
-            score += per_doc_approx_scores[k];
-            per_doc_approx_scores[k] = -9999;
-        }
-        maxsim_args->approx_scores.push(std::make_pair(score, pid));
-        seen_codes.clear();
-    }
-
-    return NULL;
-}
-
-std::vector<int> filter_pids_helper(int ncentroids, int nquery_vectors, int npids,
-                        int* pids, float* centroid_scores, int* codes,
-                        int64_t* doclens, int64_t* offsets, bool* idx,
-                        int nfiltered_docs) {
-    auto nthreads = at::get_num_threads();
-
-    pthread_t threads[nthreads];
-    maxsim_args_t args[nthreads];
-
-    for (int i = 0; i < nthreads; i++) {
-        args[i].tid = i;
-        args[i].nthreads = nthreads;
-
-        args[i].ncentroids = ncentroids;
-        args[i].nquery_vectors = nquery_vectors;
-        args[i].npids = npids;
-
-        args[i].pids = pids;
-        args[i].centroid_scores = centroid_scores;
-        args[i].codes = codes;
-        args[i].doclens = doclens;
-        args[i].offsets = offsets;
-        args[i].idx = idx;
-
-        args[i].approx_scores = std::priority_queue<std::pair<float, int>>();
-
-        int rc = pthread_create(&threads[i], NULL, maxsim, (void*)&args[i]);
-        if (rc) {
-            fprintf(stderr, "Unable to create thread %d: %d\n", i, rc);
-            std::exit(1);
-        }
-    }
-
-    for (int i = 0; i < nthreads; i++) {
-        pthread_join(threads[i], NULL);
-    }
-
-    std::priority_queue<std::pair<float, int>> global_approx_scores;
-    for (int i = 0; i < nthreads; i++) {
-        for (int j = 0; j < nfiltered_docs; j++) {
-            if (args[i].approx_scores.empty()) {
-                break;
-            }
-            global_approx_scores.push(args[i].approx_scores.top());
-            args[i].approx_scores.pop();
-        }
-    }
-
-    std::vector<int> filtered_pids;
-    for (int i = 0; i < nfiltered_docs; i++) {
-        if (global_approx_scores.empty()) {
-            break;
-        }
-        std::pair<float, int> score_and_pid = global_approx_scores.top();
-        global_approx_scores.pop();
-        filtered_pids.push_back(score_and_pid.second);
-    }
-
-    return filtered_pids;
-}
-
-torch::Tensor filter_pids(const torch::Tensor pids,
-                          const torch::Tensor centroid_scores,
-                          const torch::Tensor codes,
-                          const torch::Tensor doclens,
-                          const torch::Tensor offsets, const torch::Tensor idx,
-                          int nfiltered_docs) {
-    auto ncentroids = centroid_scores.size(0);
-    auto nquery_vectors = centroid_scores.size(1);
-    auto npids = pids.size(0);
-
-    auto pids_a = pids.data_ptr<int>();
-    auto centroid_scores_a = centroid_scores.data_ptr<float>();
-    auto codes_a = codes.data_ptr<int>();
-    auto doclens_a = doclens.data_ptr<int64_t>();
-    auto offsets_a = offsets.data_ptr<int64_t>();
-    auto idx_a = idx.data_ptr<bool>();
-
-    std::vector<int> filtered_pids = filter_pids_helper(ncentroids, nquery_vectors, npids, pids_a,
-                       centroid_scores_a, codes_a, doclens_a, offsets_a, idx_a,
-                       nfiltered_docs);
-
-    int nfinal_filtered_docs = (int)(nfiltered_docs / 4);
-    bool ones[ncentroids];
-    for (int i = 0; i < ncentroids; i++) {
-        ones[i] = true;
-    }
-
-    int* filtered_pids_a = filtered_pids.data();
-    auto nfiltered_pids = filtered_pids.size();
-    std::vector<int> final_filtered_pids = filter_pids_helper(ncentroids, nquery_vectors, nfiltered_pids,
-                       filtered_pids_a, centroid_scores_a, codes_a, doclens_a,
-                       offsets_a, ones, nfinal_filtered_docs);
-
-    auto options =
-        torch::TensorOptions().dtype(torch::kInt32).requires_grad(false);
-    return torch::from_blob(final_filtered_pids.data(), {(int)final_filtered_pids.size()},
-                            options)
-        .clone();
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("filter_pids_cpp", &filter_pids, "Filter pids", py::call_guard<py::gil_scoped_release>());
-}

src/extras/segmented_lookup.cpp

@@ -1,148 +0,0 @@
-#include <pthread.h>
-#include <torch/extension.h>
-
-#include <algorithm>
-#include <numeric>
-
-typedef struct {
-    int tid;
-    pthread_mutex_t* mutex;
-    std::queue<int>* queue;
-
-    int64_t ndocs;
-    int64_t noutputs;
-    int64_t dim;
-
-    void* input;
-    int64_t* lengths;
-    int64_t* offsets;
-    int64_t* cumulative_lengths;
-
-    void* output;
-} lookup_args_t;
-
-template <typename T>
-void* lookup(void* args) {
-    lookup_args_t* lookup_args = (lookup_args_t*)args;
-
-    int64_t* lengths = lookup_args->lengths;
-    int64_t* cumulative_lengths = lookup_args->cumulative_lengths;
-    int64_t* offsets = lookup_args->offsets;
-    int64_t dim = lookup_args->dim;
-
-    T* input = static_cast<T*>(lookup_args->input);
-    T* output = static_cast<T*>(lookup_args->output);
-
-    while (1) {
-        pthread_mutex_lock(lookup_args->mutex);
-        if (lookup_args->queue->empty()) {
-            pthread_mutex_unlock(lookup_args->mutex);
-            return NULL;
-        }
-        int i = lookup_args->queue->front();
-        lookup_args->queue->pop();
-        pthread_mutex_unlock(lookup_args->mutex);
-
-        std::memcpy(output + (cumulative_lengths[i] * dim),
-                    input + (offsets[i] * dim), lengths[i] * dim * sizeof(T));
-    }
-}
-
-template <typename T>
-torch::Tensor segmented_lookup_impl(const torch::Tensor input,
-                                    const torch::Tensor pids,
-                                    const torch::Tensor lengths,
-                                    const torch::Tensor offsets) {
-    auto lengths_a = lengths.data_ptr<int64_t>();
-    auto offsets_a = offsets.data_ptr<int64_t>();
-
-    int64_t ndocs = pids.size(0);
-    int64_t noutputs = std::accumulate(lengths_a, lengths_a + ndocs, 0);
-
-    int nthreads = at::get_num_threads();
-
-    int64_t dim;
-    torch::Tensor output;
-
-    if (input.dim() == 1) {
-        dim = 1;
-        output = torch::zeros({noutputs}, input.options());
-    } else {
-        assert(input.dim() == 2);
-        dim = input.size(1);
-        output = torch::zeros({noutputs, dim}, input.options());
-    }
-
-    int64_t cumulative_lengths[ndocs + 1];
-    cumulative_lengths[0] = 0;
-    std::partial_sum(lengths_a, lengths_a + ndocs, cumulative_lengths + 1);
-
-    pthread_mutex_t mutex;
-    int rc = pthread_mutex_init(&mutex, NULL);
-    if (rc) {
-        fprintf(stderr, "Unable to init mutex: %d\n", rc);
-    }
-
-    std::queue<int> queue;
-    for (int i = 0; i < ndocs; i++) {
-        queue.push(i);
-    }
-
-    pthread_t threads[nthreads];
-    lookup_args_t args[nthreads];
-    for (int i = 0; i < nthreads; i++) {
-        args[i].tid = i;
-        args[i].mutex = &mutex;
-        args[i].queue = &queue;
-
-        args[i].ndocs = ndocs;
-        args[i].noutputs = noutputs;
-        args[i].dim = dim;
-
-        args[i].input = (void*)input.data_ptr<T>();
-        args[i].lengths = lengths_a;
-        args[i].offsets = offsets_a;
-        args[i].cumulative_lengths = cumulative_lengths;
-
-        args[i].output = (void*)output.data_ptr<T>();
-
-        rc = pthread_create(&threads[i], NULL, lookup<T>, (void*)&args[i]);
-        if (rc) {
-            fprintf(stderr, "Unable to create thread %d: %d\n", i, rc);
-        }
-    }
-
-    for (int i = 0; i < nthreads; i++) {
-        pthread_join(threads[i], NULL);
-    }
-
-    rc = pthread_mutex_destroy(&mutex);
-    if (rc) {
-        fprintf(stderr, "Unable to destroy mutex: %d\n", rc);
-    }
-
-    return output;
-}
-
-torch::Tensor segmented_lookup(const torch::Tensor input,
-                               const torch::Tensor pids,
-                               const torch::Tensor lengths,
-                               const torch::Tensor offsets) {
-    if (input.dtype() == torch::kUInt8) {
-        return segmented_lookup_impl<uint8_t>(input, pids, lengths, offsets);
-    } else if (input.dtype() == torch::kInt32) {
-        return segmented_lookup_impl<int>(input, pids, lengths, offsets);
-    } else if (input.dtype() == torch::kInt64) {
-        return segmented_lookup_impl<int64_t>(input, pids, lengths, offsets);
-    } else if (input.dtype() == torch::kFloat32) {
-        return segmented_lookup_impl<float>(input, pids, lengths, offsets);
-    } else if (input.dtype() == torch::kFloat16) {
-        return segmented_lookup_impl<at::Half>(input, pids, lengths, offsets);
-    } else {
-        assert(false);
-    }
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("segmented_lookup_cpp", &segmented_lookup, "Segmented lookup");
-}

src/extras/segmented_maxsim.cpp

@@ -1,97 +0,0 @@
-#include <pthread.h>
-#include <torch/extension.h>
-
-#include <algorithm>
-#include <numeric>
-
-typedef struct {
-    int tid;
-    int nthreads;
-
-    int ndocs;
-    int ndoc_vectors;
-    int nquery_vectors;
-
-    int64_t* lengths;
-    float* scores;
-    int64_t* offsets;
-
-    float* max_scores;
-} max_args_t;
-
-void* max(void* args) {
-    max_args_t* max_args = (max_args_t*)args;
-
-    int ndocs_per_thread =
-        std::ceil(((float)max_args->ndocs) / max_args->nthreads);
-    int start = max_args->tid * ndocs_per_thread;
-    int end = std::min((max_args->tid + 1) * ndocs_per_thread, max_args->ndocs);
-
-    auto max_scores_offset =
-        max_args->max_scores + (start * max_args->nquery_vectors);
-    auto scores_offset =
-        max_args->scores + (max_args->offsets[start] * max_args->nquery_vectors);
-
-    for (int i = start; i < end; i++) {
-        for (int j = 0; j < max_args->lengths[i]; j++) {
-            std::transform(max_scores_offset,
-                           max_scores_offset + max_args->nquery_vectors,
-                           scores_offset, max_scores_offset,
-                           [](float a, float b) { return std::max(a, b); });
-            scores_offset += max_args->nquery_vectors;
-        }
-        max_scores_offset += max_args->nquery_vectors;
-    }
-
-    return NULL;
-}
-
-torch::Tensor segmented_maxsim(const torch::Tensor scores,
-                               const torch::Tensor lengths) {
-    auto lengths_a = lengths.data_ptr<int64_t>();
-    auto scores_a = scores.data_ptr<float>();
-    auto ndocs = lengths.size(0);
-    auto ndoc_vectors = scores.size(0);
-    auto nquery_vectors = scores.size(1);
-    auto nthreads = at::get_num_threads();
-
-    torch::Tensor max_scores =
-        torch::zeros({ndocs, nquery_vectors}, scores.options());
-
-    int64_t offsets[ndocs + 1];
-    offsets[0] = 0;
-    std::partial_sum(lengths_a, lengths_a + ndocs, offsets + 1);
-
-    pthread_t threads[nthreads];
-    max_args_t args[nthreads];
-
-    for (int i = 0; i < nthreads; i++) {
-        args[i].tid = i;
-        args[i].nthreads = nthreads;
-
-        args[i].ndocs = ndocs;
-        args[i].ndoc_vectors = ndoc_vectors;
-        args[i].nquery_vectors = nquery_vectors;
-
-        args[i].lengths = lengths_a;
-        args[i].scores = scores_a;
-        args[i].offsets = offsets;
-
-        args[i].max_scores = max_scores.data_ptr<float>();
-
-        int rc = pthread_create(&threads[i], NULL, max, (void*)&args[i]);
-        if (rc) {
-            fprintf(stderr, "Unable to create thread %d: %d\n", i, rc);
-        }
-    }
-
-    for (int i = 0; i < nthreads; i++) {
-        pthread_join(threads[i], NULL);
-    }
-
-    return max_scores.sum(1);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("segmented_maxsim_cpp", &segmented_maxsim, "Segmented MaxSim");
-}

src/index.py

@@ -1,67 +0,0 @@
-import os
-os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Prevents deadlocks in ColBERT tokenization
-os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"     # Allows multiple libraries in OpenMP runtime. This can cause unexected behavior, but allows ColBERT to work
-
-import json
-
-from constants import INDEX_NAME, DATASET_PATH
-
-from colbert import Indexer, Searcher
-from colbert.infra import Run, RunConfig, ColBERTConfig
-
-
-nbits      = 2     # encode each dimension with 2 bits
-doc_maxlen = 512   # truncate passages
-checkpoint = 'colbert-ir/colbertv2.0' # ColBERT model to use
-
-
-def index_anthology(collection, index_name):
-    with Run().context(RunConfig(nranks=2, experiment='notebook')): # nranks specifies the number of GPUs to use
-        config = ColBERTConfig(
-            doc_maxlen=doc_maxlen, 
-            nbits=nbits, 
-            kmeans_niters=4, # specifies the number of iterations of k-means clustering; 4 is a good and fast default.
-            index_path=INDEX_NAME,
-            bsize=1024
-        ) 
-
-        indexer = Indexer(
-            checkpoint=checkpoint, 
-            config=config
-        )
-
-        indexer.index(
-            name=index_name, 
-            collection=collection, 
-            overwrite=True
-        )
-
-
-def search_anthology(query, collection, index_name):
-    """ Default ColBERT search function """
-    with Run().context(RunConfig(nranks=0, experiment='notebook')):
-        searcher = Searcher(index=index_name, collection=collection)
-
-    results = searcher.search(query, k=3)
-
-    for passage_id, passage_rank, passage_score in zip(*results):
-        print(f"\t [{passage_rank}] \t\t {passage_score:.1f} \t\t {searcher.collection[passage_id]}")
-
-
-def main():
-    # Load the parsed anthology
-    with open(DATASET_PATH, 'r', encoding='utf-8') as f:
-        dataset = json.loads(f.read())
-    
-    # Get the abstracts for indexing
-    collection = [e['abstract'] for e in dataset]
-
-    # Run ColBERT indexer
-    index_anthology(collection, index_name=INDEX_NAME)
-
-    # Sanity check
-    # query = ["What are some recent examples of grammar checkers?"]
-    # search_anthology(query, collection, index_name=INDEX_NAME)
-
-
-if __name__ == '__main__': main()

src/parse.py

@@ -1,113 +0,0 @@
-import bibtexparser, json
-
-from constants import ANTHOLOGY_PATH, DATASET_PATH
-
-
-def parse_bibtex(anthology_path, dataset_path):
-    with open(anthology_path, 'r', encoding='utf-8') as f:
-        bib = bibtexparser.load(f)
-    dataset = bib.entries
-
-    print(f'Found {len(dataset)} articles with keys: {dataset[0].keys()}')
-    paper: dict
-    for paper in dataset[:2]:
-        print(f"{paper.get('author')}\n{paper.get('title')}\n{paper.get('url')}\n")
-
-    # Remove any entries without abstracts, since we index on abstracts
-    dataset = [paper for paper in dataset if 'abstract' in paper.keys()]
-
-    with open(dataset_path, 'w', encoding='utf-8') as f:
-        f.write(json.dumps(dataset, indent=4))
-
-    return dataset
-
-
-def preprocess_acl_entries(dataset_path):
-    """
-    Very rough attempt at using ACL URLs to infer their venues. Bless this mess. 
-    """
-    with open(dataset_path, 'r', encoding='utf-8') as f:
-        dataset = json.loads(f.read())
-
-    venues = []
-    for id, paper in enumerate(dataset):
-        url = paper['url']
-        year = int(paper['year'])
-
-        if year < 2020: 
-            dataset[id]['findings'] = None
-            dataset[id]['venue_type'] = None
-            continue
-
-        if 'https://aclanthology.org/' in url:
-            url = url.split('https://aclanthology.org/')[1]
-        elif 'http://www.lrec-conf.org/proceedings/' in url:
-            url = url.split('http://www.lrec-conf.org/proceedings/')[1]
-
-        if year >= 2020:
-            # new URL format
-
-            url_new = '.'.join(url.split('.')[:-1])
-            if url_new != '': url = url_new
-
-            # For most new venues, the format is "2023.eacl-tutorials" -> "eacl-tutorials"
-            url_new = '.'.join(url.split('.')[1:])
-            if url_new != '': url = url_new
-
-            # 'acl-main' -> 'acl-long'?
-            # 'acl-main' -> 'acl-short'?
-
-            # 'eacl-demo' -> 'eacl-demos'
-            # 'emnlp-tutorial' -> 'emnlp-tutorials'
-            url = url.replace('-demos', '-demo')
-            url = url.replace('-tutorials', '-tutorial')
-
-        elif year >= 2016:
-            # old URL format
-            # P17-1001 -> P17
-
-            url = url.split('-')[0]
-
-            raise RuntimeError('not working')
-
-        venues += [url]
-
-        # Extract paper type from URL
-        _type = None
-        if any(venue in url for venue in ['parlaclarin', 'nlpcovid19']):
-            _type = 'workshop'
-        elif not any(venue in url for venue in ['aacl', 'naacl', 'acl', 'emnlp', 'eacl', 'tacl']):
-            _type = 'workshop'
-        elif 'tacl' in url: _type = 'journal'
-        elif 'srw' in url: _type = 'workshop'
-        elif 'short' in url: _type = 'short'
-        elif 'demo' in url: _type = 'demo'
-        elif 'tutorial' in url: _type = 'tutorial'
-        elif 'industry' in url: _type = 'industry'
-        elif 'findings' in url: _type = 'findings'
-        elif 'main' in url or 'long' in url: _type = 'main'
-        else:
-            print(f'Could not parse: {url}')
-
-        findings = ('findings' in url)
-
-        dataset[id]['findings'] = findings
-        dataset[id]['venue_type'] = _type
-
-    # print(set(venues))
-
-    with open(DATASET_PATH, 'w', encoding='utf-8') as f:
-        f.write(json.dumps(dataset, indent=4))
-
-    return dataset
-
-
-def main():
-    # 1) Parse and save the anthology dataset
-    dataset = parse_bibtex(ANTHOLOGY_PATH, DATASET_PATH)
-
-    # 2) Pre-process the ACL anthology
-    dataset = preprocess_acl_entries(DATASET_PATH)
-
-
-if __name__ == '__main__': main()

src/search.py

@@ -1,204 +0,0 @@
-import os, ujson, tqdm
-import torch
-import torch.nn.functional as F
-
-from colbert import Checkpoint
-from colbert.infra.config import ColBERTConfig
-from colbert.search.index_storage import IndexScorer
-from colbert.search.strided_tensor import StridedTensor
-from colbert.indexing.codecs.residual_embeddings_strided import ResidualEmbeddingsStrided
-from colbert.indexing.codecs.residual import ResidualCodec
-
-
-NCELLS = 1  # Number of centroids to use in PLAID
-CENTROID_SCORE_THRESHOLD = 0.5 # How close a document has to be to a centroid to be considered
-NDOCS = 512  # Number of closest documents to consider
-
-
-def init_colbert(index_path, load_index_with_mmap=False):
-    """
-    Load all tensors necessary for running ColBERT
-    """
-    global index_checkpoint, scorer, centroids, embeddings, ivf, doclens, nbits, bucket_weights, codec, offsets
-
-    # index_checkpoint: Checkpoint    
-    
-    use_gpu = torch.cuda.is_available()
-    if use_gpu: 
-        device = 'cuda'
-    else:
-        device = 'cpu'
-
-    # Load index checkpoint
-    from colbert.infra.run import Run
-    initial_config    = ColBERTConfig.from_existing(None, Run().config)
-    index_config      = ColBERTConfig.load_from_index(index_path)
-    checkpoint_path = index_config.checkpoint
-    checkpoint_config = ColBERTConfig.load_from_checkpoint(checkpoint_path)
-    config: ColBERTConfig = ColBERTConfig.from_existing(checkpoint_config, index_config, initial_config)
-
-    index_checkpoint = Checkpoint(checkpoint_path, colbert_config=config)
-    index_checkpoint = index_checkpoint.to(device)
-    
-    load_index_with_mmap = config.load_index_with_mmap
-    if load_index_with_mmap and use_gpu:
-        raise ValueError(f"Memory-mapped index can only be used with CPU!")
-
-    scorer = IndexScorer(index_path, use_gpu, load_index_with_mmap)
-
-    with open(os.path.join(index_path, 'metadata.json')) as f:
-        metadata = ujson.load(f)
-    nbits = metadata['config']['nbits']
-
-    centroids = torch.load(os.path.join(index_path, 'centroids.pt'), map_location=device)
-    centroids = centroids.float()
-
-    ivf, ivf_lengths = torch.load(os.path.join(index_path, "ivf.pid.pt"), map_location=device)
-    ivf = StridedTensor(ivf, ivf_lengths, use_gpu=False)
-
-    embeddings = ResidualCodec.Embeddings.load_chunks(
-        index_path,
-        range(metadata['num_chunks']),
-        metadata['num_embeddings'],
-        load_index_with_mmap=load_index_with_mmap,
-    )
-
-    doclens = []
-    for chunk_idx in tqdm.tqdm(range(metadata['num_chunks'])):
-        with open(os.path.join(index_path, f'doclens.{chunk_idx}.json')) as f:
-            chunk_doclens = ujson.load(f)
-            doclens.extend(chunk_doclens)
-    doclens = torch.tensor(doclens)
-
-    buckets_path = os.path.join(index_path, 'buckets.pt')
-    bucket_cutoffs, bucket_weights = torch.load(buckets_path, map_location=device)
-    bucket_weights = bucket_weights.float()
-
-    codec = ResidualCodec.load(index_path)
-
-    if load_index_with_mmap:
-        assert metadata['num_chunks'] == 1
-        offsets = torch.cumsum(doclens, dim=0)
-        offsets = torch.cat((torch.zeros(1, dtype=torch.int64), offsets))
-    else:
-        embeddings_strided = ResidualEmbeddingsStrided(codec, embeddings, doclens)
-        offsets = embeddings_strided.codes_strided.offsets
-
-
-def colbert_score(Q: torch.Tensor, D_padded: torch.Tensor, D_mask: torch.Tensor) -> torch.Tensor:
-    """
-    Computes late interaction between question (Q) and documents (D)
-    See Figure 1: https://aclanthology.org/2022.naacl-main.272.pdf#page=3
-    """
-    assert Q.dim() == 3, Q.size()
-    assert D_padded.dim() == 3, D_padded.size()
-    assert Q.size(0) in [1, D_padded.size(0)]
-
-    scores_padded = D_padded @ Q.to(dtype=D_padded.dtype).permute(0, 2, 1)
-
-    D_padding = ~D_mask.view(scores_padded.size(0), scores_padded.size(1)).bool()
-    scores_padded[D_padding] = -9999
-    scores = scores_padded.max(1).values
-    scores = scores.sum(-1)
-
-    return scores
-
-
-def get_candidates(Q: torch.Tensor, ivf: StridedTensor) -> torch.Tensor:
-    """
-    First find centroids closest to Q, then return all the passages in all
-    centroids.
-
-    We can replace this function with a k-NN search finding the closest passages
-    using BERT similarity.
-    """
-    Q = Q.squeeze(0)
-
-    # Get the closest centroids via a matrix multiplication + argmax
-    centroid_scores: torch.Tensor = (centroids @ Q.T)
-    if NCELLS == 1:
-        cells = centroid_scores.argmax(dim=0, keepdim=True).permute(1, 0)
-    else:
-        cells = centroid_scores.topk(NCELLS, dim=0, sorted=False).indices.permute(1, 0)  # (32, ncells)
-    cells = cells.flatten().contiguous()  # (32 * ncells,)
-    cells = cells.unique(sorted=False)
-
-    # Given the relevant clusters, get all passage IDs in each cluster
-    # Note, this may return duplicates since passages can exist in multiple clusters
-    pids, _ = ivf.lookup(cells)
-
-    # Sort and retun values
-    pids = pids.sort().values
-    pids, _ = torch.unique_consecutive(pids, return_counts=True)
-    return pids, centroid_scores
-
-
-def _calculate_colbert(Q: torch.Tensor):
-    """
-    Multi-stage ColBERT pipeline. Implemented using the PLAID engine, see fig. 5:
-    https://arxiv.org/pdf/2205.09707#page=5
-    """
-    # Stage 1 (Initial Candidate Generation): Find the closest candidates to the Q centroid score
-    unfiltered_pids, centroid_scores = get_candidates(Q, ivf)
-    print(f'Stage 1 candidate generation: {unfiltered_pids.shape}')
-
-    # Stage 2 and 3 (Centroid Interaction with Pruning, then without Pruning)
-    idx = centroid_scores.max(-1).values >= CENTROID_SCORE_THRESHOLD
-
-    # C++ : Filter pids under the centroid score threshold
-    pids_true = scorer.filter_pids(
-        unfiltered_pids, centroid_scores, embeddings.codes, doclens, offsets, idx, NDOCS
-    )
-    pids = pids_true
-    assert torch.equal(pids_true, pids), f'\n{pids_true}\n{pids}'
-    print('Stage 2 filtering:', unfiltered_pids.shape, '->', pids.shape) # (n_docs) -> (n_docs/4)
-
-    # Stage 3.5 (Decompression) - Get the true passage embeddings for calculating maxsim
-    D_packed = scorer.decompress_residuals(
-        pids, doclens, offsets, bucket_weights, codec.reversed_bit_map,
-        codec.decompression_lookup_table, embeddings.residuals, embeddings.codes, 
-        centroids, codec.dim, nbits
-    )
-    D_packed = F.normalize(D_packed.to(torch.float32), p=2, dim=-1)
-    D_mask = doclens[pids.long()]
-    D_padded, D_lengths = StridedTensor(D_packed, D_mask, use_gpu=False).as_padded_tensor()
-    print('Stage 3.5 decompression:', pids.shape, '->', D_padded.shape) # (n_docs/4) -> (n_docs/4, num_toks, hidden_dim)
-
-    # Stage 4 (Final Ranking w/ Decompression) - Calculate the final (expensive) maxsim scores with ColBERT 
-    scores = colbert_score(Q, D_padded, D_lengths)
-    print('Stage 4 ranking:', D_padded.shape, '->', scores.shape)
-
-    return scores, pids
-
-
-def encode(text, full_length_search=False) -> torch.Tensor:
-    queries = text if isinstance(text, list) else [text]
-    bsize = 128 if len(queries) > 128 else None
-
-    Q = index_checkpoint.queryFromText(
-        queries, 
-        bsize=bsize, 
-        to_cpu=True, 
-        full_length_search=full_length_search
-    )
-
-    QUERY_MAX_LEN = index_checkpoint.query_tokenizer.query_maxlen
-    Q = Q[:, :QUERY_MAX_LEN] # Cut off query to maxlen tokens
-    
-    return Q
-
-
-def search_colbert(query):
-    """
-    ColBERT search with a query.
-    """
-    # Encode query using ColBERT model, using the appropriate [Q], [D] tokens
-    Q = encode(query)    
-
-    scores, pids = _calculate_colbert(Q)
-
-    # Sort values
-    scores_sorter = scores.sort(descending=True)
-    pids, scores = pids[scores_sorter.indices].tolist(), scores_sorter.values.tolist()
-
-    return pids, scores

src/server.py

@@ -1,106 +0,0 @@
-import os, math, re
-from typing import List, Optional, Union
-
-from flask import Flask, abort, request, render_template
-from functools import lru_cache
-
-from constants import INDEX_PATH, VENUES
-
-from search import init_colbert, search_colbert
-from db import create_database, query_paper_metadata
-
-PORT = int(os.getenv("PORT", 8893))
-app = Flask(__name__)
-
-
-@lru_cache(maxsize=1000000)
-def api_search_query(query):
-    print(f"Query={query}")
-
-    # Use ColBERT to find passages related to the query
-    pids, scores = search_colbert(query)
-
-    # Softmax output probs
-    probs = [math.exp(s) for s in scores]
-    probs = [p / sum(probs) for p in probs]
-
-    # Sort and return results as a dict
-    topk = [{'pid': pid, 'score': score, 'prob': prob} for pid, score, prob in zip(pids, scores, probs)]
-    topk = sorted(topk, key=lambda p: (p['score'], p['pid']), reverse=True)
-
-    response = {"query" : query, "topk": topk}
-
-    return response
-
-
-def is_valid_query(query):
-    return re.match(r'^[a-zA-Z0-9 ]*$', query) and len(query) <= 256
-
-
-@app.route("/api/colbert", methods=["GET"])
-def api_search():
-    if request.method == "GET":
-        query = str(request.args.get('query'))
-        if not is_valid_query(query): abort(400, "Invalid query :(")
-        return api_search_query(query)
-    return ('', 405)
-
-
-@app.route('/api/search', methods=['POST', 'GET'])
-def query():
-    query: str
-    start_year: Optional[int]
-    end_year: Optional[int]
-    venue_type: Optional[Union[VENUES, List[VENUES]]]
-    is_findings: Optional[bool]
-
-    if request.method in ["POST", "GET"]:
-        args = request.form if request.method == "POST" else request.args
-        query       = args.get('query')
-        start_year  = args.get('start_year', None)
-        end_year    = args.get('end_year', None)
-        venue_type  = args.getlist('venue_type', None)
-        is_findings = args.get('is_findings', None)
-    
-    if not is_valid_query(query): 
-        abort(400, "Invalid query :(")
-
-    # Get top passage IDs from ColBERT
-    colbert_response = api_search_query(query)
-
-    # Query MySQL database for paper information
-    pids = [r['pid'] for r in colbert_response["topk"]]
-    mysql_response = query_paper_metadata(
-        pids, 
-        start_year=start_year,
-        end_year=end_year,
-        venue_type=venue_type,
-        is_findings=is_findings
-    )
-
-    K = 20
-    mysql_response = mysql_response[:K]
-
-    return mysql_response
-
-
-# @app.route('/search', methods=['POST', 'GET'])
-# def search_web():
-#     return render_template('public/results.html', query=query, year=year, results=results)
-    
-
-@app.route('/', methods=['POST', 'GET'])
-def index():
-    return render_template('index.html')
-
-
-if __name__ == "__main__":
-    """
-    Example usage:
-    python server.py
-    http://localhost:8893/api/colbert?query=Information retrevial with BERT
-    http://localhost:8893/api/search?query=Information retrevial with BERT
-    """
-    create_database()
-    init_colbert(index_path=INDEX_PATH)
-    app.run("0.0.0.0", PORT) # debug=True

src/static/style.css

@@ -1,83 +0,0 @@
-:root {
-    --custom-red: #ED1C24;
-    --custom-red-dark: #D11920;
-    --custom-blue: #446e9b;
-}
-.card {
-    margin-bottom: 20px;
-}
-.search-container {
-    display: flex;
-    margin-bottom: 20px;
-}
-.search-container .form-control {
-    margin-right: 10px;
-}
-.btn-primary {
-    background-color: var(--custom-red);
-    border-color: var(--custom-red);
-}
-.btn-primary:hover, .btn-primary:focus, .btn-primary:active {
-    background-color: var(--custom-red-dark);
-    border-color: var(--custom-red-dark);
-}
-/* Custom styling for range input */
-input[type="range"] {
-    width: 100%;
-    height: 8px;
-    border-radius: 5px;
-    background: #b6b6b6;
-    outline: none;
-}
-input[type="range"]::-webkit-slider-thumb {
-    -webkit-appearance: none;
-    appearance: none;
-    width: 20px;
-    height: 20px;
-    border-radius: 50%;
-    background: var(--custom-red);
-    cursor: pointer;
-}
-input[type="range"]::-moz-range-thumb {
-    width: 20px;
-    height: 20px;
-    border-radius: 50%;
-    background: var(--custom-red);
-    cursor: pointer;
-}
-/* Custom styling for checkboxes */
-.form-check-input:checked {
-    background-color: var(--custom-red);
-    border-color: var(--custom-red);
-}
-.form-check-input:focus {
-    border-color: var(--custom-red);
-    box-shadow: 0 0 0 0.25rem rgba(237, 28, 36, 0.25);
-}
-/* Custom styling for links */
-a {
-    color: var(--custom-blue) !important;
-    text-decoration: none;
-}
-a:hover, a:focus {
-    color: #446e9b !important;
-    text-decoration: underline;
-}
-/* David custom */
-.card-text:last-child {
-    font-size: 11pt;
-    line-height: 1.05 !important;
-}
-.card {
-    background-color: #f8f9fa !important;
-    color: #212529
-}
-.card-body h6 {
-    font-size: 11pt;
-}
-.paper-metadata {
-    font-size: 10pt;
-}
-.range-label {
-    text-align: center
-}

src/templates/index.html

@@ -1,137 +0,0 @@
-<!DOCTYPE html>
-<html lang="en">
-
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>ACL Search</title>
-    <link href="https://cdnjs.cloudflare.com/ajax/libs/bootstrap/5.3.0/css/bootstrap.min.css" rel="stylesheet">
-    <link rel="stylesheet" href="{{ url_for('static', filename='style.css') }}">
-    <script src="https://cdnjs.cloudflare.com/ajax/libs/axios/0.21.1/axios.min.js"></script>
-</head>
-
-<body>
-    <div class="container mt-5">
-        <form id="searchForm">
-            <div class="row">
-                <div class="col-md-3">
-                    <div class="mb-3">
-                        <h1 class="mb-4">ACL Search</h1>
-                        <label for="yearRange" class="form-label range-label"><span id="yearRangeValue"></span></label>
-                        <input type="range" class="form-range" id="yearRange" min="2010" max="2024" step="1"
-                            value="2021">
-                    </div>
-
-                    <div class="mb-3">
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="main" id="main" checked>
-                            <label class="form-check-label" for="main">Main Long</label>
-                        </div>
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="short" id="short" checked>
-                            <label class="form-check-label" for="short">Main Short</label>
-                        </div>
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="findings" id="findings">
-                            <label class="form-check-label" for="findings">Findings</label>
-                        </div>
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="journal" id="journal">
-                            <label class="form-check-label" for="journal">Journal</label>
-                        </div>
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="workshop" id="workshop">
-                            <label class="form-check-label" for="workshop">Workshop</label>
-                        </div>
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="demo" id="demo">
-                            <label class="form-check-label" for="demo">Demo Track</label>
-                        </div>
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="industry" id="industry">
-                            <label class="form-check-label" for="industry">Industry Track</label>
-                        </div>
-                        <div class="form-check">
-                            <input class="form-check-input" type="checkbox" value="tutorial" id="tutorial">
-                            <label class="form-check-label" for="tutorial">Tutorial Abstracts</label>
-                        </div>
-                    </div>
-                </div>
-                <div class="col-md-9">
-                    <div class="search-container">
-                        <input type="text" class="form-control" id="query" name="query" placeholder="Information is the resolution of uncertainty" required>
-                        <button type="submit" class="btn btn-primary" form="searchForm">Search</button>
-                    </div>
-                    <div id="results" class="mt-4"></div>
-                </div>
-            </div>
-        </form>
-    </div>
-
-    <script>
-        const yearRange = document.getElementById('yearRange');
-        const yearRangeValue = document.getElementById('yearRangeValue');
-
-        function updateYearRangeValue() {
-            const startYear = parseInt(yearRange.value);
-            const endYear = 2024;
-            yearRangeValue.textContent = `${startYear} - ${endYear}`;
-        }
-
-        yearRange.addEventListener('input', updateYearRangeValue);
-        updateYearRangeValue();  // Initial call to set the text
-
-        document.getElementById('searchForm').addEventListener('submit', async (e) => {
-            e.preventDefault();
-            const form = e.target;
-
-            const params = new URLSearchParams();
-            params.append('query', form.query.value);
-            params.append('start_year', yearRange.value);
-            params.append('end_year', '2024');
-
-            const selectedVenues = Array.from(document.querySelectorAll('input[type="checkbox"]:checked'))
-                .map(checkbox => checkbox.value);
-            selectedVenues.forEach(venue => params.append('venue_type', venue));
-
-            try {
-                const response = await axios.get('http://localhost:8893/api/search', { params });
-                displayResults(response.data);
-            } catch (error) {
-                console.error('Error:', error);
-                document.getElementById('results').innerHTML = '<p class="alert alert-danger">An error occurred while fetching results.</p>';
-            }
-        });
-
-        function displayResults(data) {
-            const resultsDiv = document.getElementById('results');
-            if (data.length === 0) {
-                resultsDiv.innerHTML = '<p class="alert alert-info">No results found.</p>';
-                return;
-            }
-
-            console.log(data)
-
-            let html = '';
-            data.forEach(paper => {
-                html += `
-                    <div class="card">
-                        <div class="card-body">
-                            <h5 class="card-title"><a href="${paper.url}" target="_blank">${paper.title}</a></h5>
-                            <h6 class="card-subtitle mb-1 text-muted">${paper.authors}</h6>
-                            <p class="card-subtitle text-muted paper-metadata">
-                                <strong>${paper.year} / ${paper.venue_type}</strong> <br>
-                                ${paper.is_findings ? '<br><strong>Findings Paper</strong>' : ''}
-                            </p>
-                            <p class="card-text"><small class="text-muted">${paper.abstract}</small></p>
-                        </div>
-                    </div>
-                `;
-            });
-
-            resultsDiv.innerHTML = html;
-        }
-    </script>
-</body>
-
-</html>

src/utils.py

@@ -1,95 +0,0 @@
-import torch
-import tqdm
-
-def maxsim(pids, centroid_scores, codes, doclens, offsets, idx, nfiltered_docs):
-    ncentroids, nquery_vectors = centroid_scores.shape
-    centroid_scores = centroid_scores.flatten()
-    scores = []
-
-    for i in tqdm.tqdm(range(len(pids)), desc='Calculating maxsim over centroids...'):
-        seen_codes = set()
-        per_doc_scores = torch.full((nquery_vectors,), -9999, dtype=torch.float32)
-
-        pid = pids[i]
-        for j in range(doclens[pid]):
-            code = codes[offsets[pid] + j]
-            assert code < ncentroids
-            if idx[code] and code not in seen_codes:
-                for k in range(nquery_vectors):
-                    per_doc_scores[k] = torch.max(
-                        per_doc_scores[k], 
-                        centroid_scores[code * nquery_vectors + k]
-                    )
-                seen_codes.add(code)
-
-        score = torch.sum(per_doc_scores[:nquery_vectors]).item()
-        scores += [(score, pid)]
-
-    # Sort and return scores
-    global_scores = sorted(scores, key=lambda x: x[0], reverse=True)
-    filtered_pids = [pid for _, pid in global_scores[:nfiltered_docs]]
-    filtered_pids = torch.tensor(filtered_pids, dtype=torch.int32)
-
-    return filtered_pids
-
-
-def filter_pids(pids, centroid_scores, codes, doclens, offsets, idx, nfiltered_docs):
-    filtered_pids = maxsim(
-        pids, centroid_scores, codes, doclens, offsets, idx, nfiltered_docs
-    )
-
-    print('Stage 2 filtering:', pids.shape, '->', filtered_pids.shape) # (all_docs) -> (n_docs/4)
-
-    nfinal_filtered_docs = int(nfiltered_docs / 4)
-    ones = [True] * centroid_scores.size(0)
-
-    final_filtered_pids = maxsim(
-        filtered_pids, centroid_scores, codes, doclens, offsets, ones, nfinal_filtered_docs
-    )
-
-    print('Stage 3 filtering:', filtered_pids.shape, '->', final_filtered_pids.shape) # (n_docs) -> (n_docs/4)
-
-    return final_filtered_pids
-
-
-def decompress_residuals(pids, doclens, offsets, bucket_weights, reversed_bit_map,
-        bucket_weight_combinations, binary_residuals, codes, 
-        centroids, dim, nbits):
-    npacked_vals_per_byte = 8 // nbits
-    packed_dim = dim // npacked_vals_per_byte
-    cumulative_lengths = [0 for _ in range(len(pids)+1)]
-    noutputs = 0
-    for i in range(len(pids)):
-        noutputs += doclens[pids[i]]
-        cumulative_lengths[i + 1] = cumulative_lengths[i] + doclens[pids[i]]
-
-    output = []
-
-    binary_residuals = binary_residuals.flatten()
-    centroids = centroids.flatten()
-
-    # Iterate over all documents
-    for i in range(len(pids)):
-        pid = pids[i]
-
-        # Offset into packed list of token vectors for the given document
-        offset = offsets[pid]
-
-        # For each document, iterate over all token vectors
-        for j in range(doclens[pid]):
-            code = codes[offset + j]
-
-            # For each token vector, iterate over the packed (8-bit) residual values
-            for k in range(packed_dim):
-                x = binary_residuals[(offset + j) * packed_dim + k]
-                x = reversed_bit_map[x]
-
-                # For each packed residual value, iterate over the bucket weight indices.
-                # If we use n-bit compression, that means there will be (8 / n) indices per packed value.
-                for l in range(npacked_vals_per_byte):
-                    output_dim_idx = k * npacked_vals_per_byte + l
-                    bucket_weight_idx = bucket_weight_combinations[x * npacked_vals_per_byte + l]
-                    output[(cumulative_lengths[i] + j) * dim + output_dim_idx] = \
-                        bucket_weights[bucket_weight_idx] + centroids[code * dim + output_dim_idx]
-
-    return output