Askem xdd-COVID retrieval-augmented generation prototype

Repo: https://github.com/AFIDSI/askem

Demo: http://cosmos0001.chtc.wisc.edu:8501/

API Base URL: http://cosmos0001.chtc.wisc.edu:4502/

For end-users

The end users of our system are ASKEM performers who access it using REST API. You can also visit our demo to try how this system can power a traceable COVID-19 search engine.

Retriever overview

The retriever uses an embedding-based search engine, specifically Dense Passage Retriever (DPR), to query relevant documents from the XDD database. Currently, it returns paragraphs as documents. Future updates may include figures, tables, and equations. The API accepts POST requests and requires an APIKEY. ASKEM performers can obtain an API key by contacting me.

Base URL: http://cosmos0001.chtc.wisc.edu:4502

There are 3 endpoints available:

1. vector: Basic DPR vector search (Not recommended).
2. hybrid: Combines Elasticsearch pre-filtering with DPR vector search (Recommended, better performance).
3. react: Builds on the hybrid approach, integrating the ReAct agent for "reasoning" (via gpt-4 by default) and subsequent querying (via hybrid endpoint by default) to generate better answers. (Experimental, slow, highest performance).

vector and hybrid endpoint example usage

Both vector and hybrid endpoints use similar format for request and response data.

import requests

APIKEY = "insert_api_key_here"
ENDPOINT = "http://cosmos0001.chtc.wisc.edu:4502/hybrid"

headers = {"Content-Type": "application/json", "Api-Key": APIKEY}
data = {
    "question": "What is SIDARTHE model?",
    "top_k": 3,
}

response = requests.post(ENDPOINT, headers=headers, json=data)
response.json()

Request body schema for vector and hybrid endpoints

{
    "question": str,
    "top_k": Optional[int] = 5, # Number of documents to return
    "distance": Optional[float] = 0.5, # Max cosine distance between question and document
    "topic": Optional[str] = None, # Filter by topic, only "covid" is available now
    "doc_type": Optional[str] = None,  # Filter by document type, only "paragraph" is available now
    "preprocessor_id": Optional[str] = None,  # Filter by preprocessor_id, for developer use only
    "article_terms": Optional[List[str]] = None,  # Obsolete, do not use
    "paragraph_terms": Optional[List[str]] = None,  # Filter by capitalized terms (any word that has more than one capital letter) in the paragraph
    "paper_ids": Optional[List[str]] = None,  # Filter by XDD paper ids
    "move_to": Optional[str] = None,  # Move the answer to better match the context of the given string, like `mathematical equation`.
    "move_to_weight": Optional[float] = 0,  # Weight `move_to` parameter to adjusts the influence on the original answer, with a range from 0 to 1. Higher values mean stronger augmentation.
    "move_away_from": Optional[str] = None,  # Move the answer away from irrelevant topics, like `general commentary`.
    "move_away_from_weight": Optional[float] = 0,  # Weight `move_away_from` to adjusts the influence on the original answer, with a range from 0 to 1. Higher values mean stronger augmentation.
    "screening_top_k": Optional[int] = 100,  # `hybrid` endpoint only. Number of documents to return from the elastic search pre-filtering step.
}

Response body schema for vector and hybrid endpoints

[
    {
        "paper_id": str,  # XDD paper id
        "doc_type": str,  # only "paragraph" for now
        "text": str,  # text content
        "distance": float,  # distance to question
        "cosmos_object_id": str,  # only available for doc_type="figure"
        "article_terms": List[str],  # Obsolete, do not use
        "paragraph_terms": List[str], # Capitalized terms in the paragraph
    },
    ...
]

react endpoint example usage

import requests

APIKEY = "insert_api_key_here"
ENDPOINT = "http://cosmos0001.chtc.wisc.edu:4502/react"

headers = {"Content-Type": "application/json", "Api-Key": APIKEY}
data = {
    "question": "What is SIDARTHE model?",
    "top_k": 3,
}

response = requests.post(ENDPOINT, headers=headers, json=data)
response.json()

Request body schema for react endpoint

{
    "question": str,
    ..., # Same as `hybrid` endpoint, see
    "retriever_endpoint": Optional[str] = "http://retriever:4502/hybrid",  # retriever endpoint to use
    "model_name": Optional[str] = "gpt-4",  # OpenAI llm model name
}

Response body schema for react endpoint

{
    "answer": str,  # Final answer to the question
    "used_docs": list[Document],  # Relevant documents used to generate the answer, with the same schema as the response of `hybrid` endpoint
}

For developer

To deploy the system

1. Make a .env file in the project root directory with these variables

   see example: .env.example

   see shared dotenv file for the actual values

2. Run launch test

   bash ./scripts/launch_test.sh

3. Ingest documents

   Put all text files in a folder, with file format as <ingest_dir>/<paper-id>.txt, then run this:

   python askem/ingest_docs.py --input-dir "data/debug_data/paragraph_test" --topic "covid-19" --doc-type "paragraph" --weaviate-url "url_to_weaviate"

4. Ingest figures

   Put all text files in a folder, with file format as <ingest_dir>/<paper-id>.<cosmos_object_id>.txt, then run this:

   python askem/deploy.py --input-dir "data/debug_data/figure_test" --topic "covid-19" --doc-type "figure" --weaviate-url "url_to_weaviate"