search

Dicklesworthstone / swiss_army_llama

A FastAPI service for semantic text search using precomputed embeddings and advanced similarity measures, with built-in support for various file types through textract.
945 stars 53 forks source link
embedding-similarity embedding-vectors embeddings llama2 llamacpp semantic-search

readme

🇨🇭🎖️🦙 Swiss Army Llama

Introduction

The Swiss Army Llama is designed to facilitate and optimize the process of working with local LLMs by using FastAPI to expose convenient REST endpoints for various tasks, including obtaining text embeddings and completions using different LLMs via llama_cpp, as well as automating the process of obtaining all the embeddings for most common document types, including PDFs (even ones that require OCR), Word files, etc; it even allows you to submit an audio file and automatically transcribes it with the Whisper model, cleans up the resulting text, and then computes the embeddings for it. To avoid wasting computation, these embeddings are cached in SQlite and retrieved if they have already been computed before. To speed up the process of loading multiple LLMs, optional RAM Disks can be used, and the process for creating and managing them is handled automatically for you. With a quick and easy setup process, you will immediately get access to a veritable "Swiss Army Knife" of LLM related tools, all accessible via a convenient Swagger UI and ready to be integrated into your own applications with minimal fuss or configuration required.

Some additional useful endpoints are provided, such as computing semantic similarity between submitted text strings. The service leverages a high-performance Rust-based library, fast_vector_similarity, to offer a range of similarity measures including spearman_rho, kendall_tau, approximate_distance_correlation, jensen_shannon_dependency_measure, and hoeffding_d. Additionally, semantic search across all your cached embeddings is supported using FAISS vector searching. You can either use the built in cosine similarity from FAISS, or supplement this with a second pass that computes the more sophisticated similarity measures for the most relevant subset of the stored vectors found using cosine similarity (see the advanced semantic search endpoint for this functionality).

Also, we now support multiple embedding pooling methods for combining token-level embedding vectors into a single fixed-length embedding vector for any length of input text, including the following:

As mentioned above, you can now submit not only plaintext and fully digital PDFs but also MS Word documents, images, and other file types supported by the textract library. The library can automatically apply OCR using Tesseract for scanned text. The returned embeddings for each sentence in a document can be organized in various formats like records, table, etc., using the Pandas to_json() function. The results can be returned either as a ZIP file containing a JSON file or as a direct JSON response. You can now also submit audio files in MP3 or WAV formats. The library uses OpenAI's Whisper model, as optimized by the Faster Whisper Python library, to transcribe the audio into text. Optionally, this transcript can be treated like any other document, with each sentence's embeddings computed and stored. The results are returned as a URL to a downloadable ZIP file containing a JSON with the embedding vector data.

Finally, we add a new endpoint for generating multiple text completions for a given input prompt, with the ability to specify a grammar file that will enforce a particular form of response, such as JSON. There is also a useful new utility feature: a real-time application log viewer that can be accessed via a web browser, which allows for syntax highlighting and offers options for downloading the logs or copying them to the clipboard. This allows a user to watch the logs without having direct SSH access to the server.

Screenshots

Swiss Army Llama Swagger UI Swiss Army Llama Runnig

TLDR: If you just want to try it very quickly on a fresh Ubuntu 22+ machine (warning, this will install docker using apt):

git clone https://github.com/Dicklesworthstone/swiss_army_llama
cd swiss_army_llama
chmod +x setup_dockerized_app_on_fresh_machine.sh
sudo ./setup_dockerized_app_on_fresh_machine.sh

To run it natively (not using Docker) in a Python venv (recommended!), you can use these commands:

sudo apt-get update
sudo apt-get install build-essential libxml2-dev libxslt1-dev antiword unrtf poppler-utils pstotext tesseract-ocr flac ffmpeg lame libmad0 libsox-fmt-mp3 sox libjpeg-dev swig redis-server libpoppler-cpp-dev pkg-config -y
sudo systemctl enable redis-server
sudo systemctl start redis
git clone https://github.com/Dicklesworthstone/swiss_army_llama
cd swiss_army_llama
python3 -m venv venv
source venv/bin/activate
python3 -m pip install --upgrade pip
python3 -m pip install wheel
python3 -m pip install --upgrade setuptools wheel
pip install -r requirements.txt
python3 swiss_army_llama.py

Alternatively, you can also just run the included script, which will install PyEnv if it's not already installed on your machine, and then install Python 3.12 and create a virtual environment for you. You can do everything with a single one-liner from scratch on a fresh Ubuntu machine like this:

git clone https://github.com/Dicklesworthstone/swiss_army_llama && cd swiss_army_llama && chmod +x install_swiss_army_llama.sh && ./install_swiss_army_llama.sh && pyenv local 3.12 && source venv/bin/activate && python swiss_army_llama.py

Then open a browser to <your_static_ip_address>:8089 if you're using a VPS to get to the FastAPI Swagger page at http://localhost:8089.

Or to localhost:8089 if you're using your own machine-- but, really, you should never run untrusted code with sudo on your own machine! Just get a cheap VPS to experiment with for $30/month.

Watch the the automated setup process in action here.

Features

  1. Text Embedding Computation: Utilizes pre-trained LLama3 and other LLMs via llama_cpp to generate embeddings for any provided text.
  2. Embedding Caching: Efficiently stores and retrieves computed embeddings in SQLite, minimizing redundant computations.
  3. Advanced Similarity Measurements and Retrieval: Utilizes the author's own fast_vector_similarity library written in Rust to offer highly optimized advanced similarity measures such as spearman_rho, kendall_tau, approximate_distance_correlation, jensen_shannon_dependency_measure, and hoeffding_d. Semantic search across cached embeddings is also supported using FAISS vector searching.
  4. Two-Step Advanced Semantic Search: The API first leverages FAISS and cosine similarity for rapid filtering, and then applies additional similarity measures like spearman_rho, kendall_tau, approximate_distance_correlation, jensen_shannon_dependency_measure, and hoeffding_d for a more nuanced comparison.
  5. File Processing for Documents: The library now accepts a broader range of file types including plaintext, PDFs, MS Word documents, and images. It can also handle OCR automatically. Returned embeddings for each sentence are organized in various formats like records, table, etc., using Pandas to_json() function.
  6. Advanced Text Preprocessing: The library now employs a more advanced sentence splitter to segment text into meaningful sentences. It handles cases where periods are used in abbreviations, domain names, or numbers and also ensures complete sentences even when quotes are used. It also takes care of pagination issues commonly found in scanned documents, such as awkward newlines and hyphenated line breaks.
  7. Audio Transcription and Embedding: Upload an audio file in MP3 or WAV format. The library uses OpenAI's Whisper model for transcription. Optionally, sentence embeddings can be computed for the transcript.
  8. RAM Disk Usage: Optionally uses RAM Disk to store models for faster access and execution. Automatically handles the creation and management of RAM Disks.
  9. Robust Exception Handling: Features comprehensive exception management to ensure system resilience.
  10. Interactive API Documentation: Integrates with Swagger UI for an interactive and user-friendly experience, accommodating large result sets without crashing.
  11. Scalability and Concurrency: Built on the FastAPI framework, handles concurrent requests and supports parallel inference with configurable concurrency levels.
  12. Flexible Configurations: Offers configurable settings through environment variables and input parameters, including response formats like JSON or ZIP files.
  13. Comprehensive Logging: Captures essential information with detailed logs, without overwhelming storage or readability.
  14. Support for Multiple Models and Measures: Accommodates multiple embedding models and similarity measures, allowing flexibility and customization based on user needs.
  15. Ability to Generate Multiple Completions using Specified Grammar: Get back structured LLM completions for a specified input prompt.
  16. Real-Time Log File Viewer in Browser: Lets anyone with access to the API server conveniently watch the application logs to gain insight into the execution of their requests.
  17. Uses Redis for Request Locking: Uses Redis to allow for multiple Uvicorn workers to run in parallel without conflicting with each other.

Demo Screen Recording in Action

Here is the live console output while I interact with it from the Swagger page to make requests.

Requirements

System requirements for running the application (to support all the file types handled by textract):

sudo apt-get update
sudo apt-get install libxml2-dev libxslt1-dev antiword unrtf poppler-utils pstotext tesseract-ocr flac ffmpeg lame libmad0 libsox-fmt-mp3 sox libjpeg-dev swig -y

Python Requirements:

aioredis
aioredlock
aiosqlite
apscheduler
faiss-cpu
fast_vector_similarity
fastapi
faster-whisper
filelock
httpx
llama-cpp-python
magika
mutagen
nvgpu
pandas
pillow
psutil
pydantic
PyPDF2
pytest
python-decouple
python-multipart
pytz
redis
ruff
scikit-learn
scipy
sqlalchemy
textract-py3
uvicorn
uvloop
zstandard

Running the Application

You can run the application using the following command:

python swiss_army_llama.py

The server will start on 0.0.0.0 at the port defined by the SWISS_ARMY_LLAMA_SERVER_LISTEN_PORT variable.

Access the Swagger UI:

http://localhost:<SWISS_ARMY_LLAMA_SERVER_LISTEN_PORT>

Configuration

You can configure the service easily by editing the included .env file. Here's a list of available configuration options:

Contributing

If you'd like to contribute to the project, please submit a pull request! Seriously, I'd love to get some more community going so we can make this a standard library!

License

This project is licensed under the MIT License.

Some Llama Knife Images I found on Google

Setup and Configuration

RAM Disk Configuration

To enable password-less sudo for RAM Disk setup and teardown, edit the sudoers file with sudo visudo. Add the following lines, replacing username with your actual username:

username ALL=(ALL) NOPASSWD: /bin/mount -t tmpfs -o size=*G tmpfs /mnt/ramdisk
username ALL=(ALL) NOPASSWD: /bin/umount /mnt/ramdisk

The application provides functionalities to set up, clear, and manage RAM Disk. RAM Disk is used to store models in memory for faster access. It calculates the available RAM and sets up the RAM Disk accordingly. The functions setup_ramdisk, copy_models_to_ramdisk, and clear_ramdisk manage these tasks.

API Endpoints

The following endpoints are available:

For detailed request and response schemas, please refer to the Swagger UI available at the root URL or the section at the end of this README.

Exception Handling

The application has robust exception handling to deal with various types of errors, including database errors and general exceptions. Custom exception handlers are defined for SQLAlchemyError and general Exception.

Logging

Logging is configured at the INFO level to provide detailed logs for debugging and monitoring. The logger provides information about the state of the application, errors, and activities.

The logs are stored in a file named swiss_army_llama.log, and a log rotation mechanism is implemented to handle log file backups. The rotating file handler is configured with a maximum file size of 10 MB, and it keeps up to 5 backup files.

When a log file reaches its maximum size, it is moved to the old_logs directory, and a new log file is created. The log entries are also printed to the standard output stream.

Here are some details of the logging configuration:

Additionally, the log level for SQLAlchemy's engine is set to WARNING to suppress verbose database logs.

Database Structure

The application uses a SQLite database via SQLAlchemy ORM. Here are the data models used, which can be found in the embeddings_data_models.py file:

TextEmbedding Table

DocumentEmbedding Table

Document Table

AudioTranscript Table

Database Relationships

  1. TextEmbedding - DocumentEmbedding:

    • TextEmbedding has a Foreign Key document_file_hash that references DocumentEmbedding's document_file_hash.
    • This means multiple text embeddings can belong to a single document embedding, establishing a one-to-many relationship.

  2. DocumentEmbedding - Document:

    • DocumentEmbedding has a Foreign Key document_hash that references Document's document_hash.
    • This establishes a one-to-many relationship between Document and DocumentEmbedding.

  3. AudioTranscript:

    • This table doesn't have a direct relationship with other tables based on the given code.

  4. Request/Response Models:

    • These are not directly related to the database tables but are used for handling API requests and responses.
    • The following Pydantic models are used for request and response validation:
      • EmbeddingRequest
      • SimilarityRequest
      • SemanticSearchRequest
      • SemanticSearchResponse
      • AdvancedSemanticSearchRequest
      • AdvancedSemanticSearchResponse
      • EmbeddingResponse
      • SimilarityResponse
      • AllStringsResponse
      • AllDocumentsResponse
      • TextCompletionRequest
      • TextCompletionResponse
      • ImageQuestionResponse
      • AudioTranscriptResponse
      • ShowLogsIncrementalModel
      • AddGrammarRequest
      • AddGrammarResponse

For detailed field descriptions and validations, please refer to the embeddings_data_models.py file.

Performance Optimizations

This section highlights the major performance enhancements integrated into the provided code to ensure swift responses and optimal resource management.

1. Asynchronous Programming:

2. Database Optimizations:

3. RAM Disk Utilization:

4. Model Caching:

5. Parallel Inference:

6. Embedding Caching:

Dockerized Version

A bash script is included in this repo, setup_dockerized_app_on_fresh_machine.sh, that will automatically do everything for you, including installing docker with apt install.

To use it, first make the script executable and then run it like this:

chmod +x setup_dockerized_app_on_fresh_machine.sh
sudo ./setup_dockerized_app_on_fresh_machine.sh

If you prefer a manual setup, then read the following instructions:

Prerequisites

Ensure that you have Docker installed on your system. If not, follow these steps to install Docker on Ubuntu:

sudo apt-get update
sudo apt-get install docker.io
sudo systemctl start docker
sudo docker --version
sudo usermod -aG docker $USER

You may need to log out and log back in or restart your system to apply the new group permissions, or use sudo in the following steps to build and run the container.

Setup and Running the Application

  1. Clone the Repository:

    Clone the Swiss Army Llama repository to your local machine:

    git clone https://github.com/Dicklesworthstone/swiss_army_llama
cd swiss_army_llama

  2. Build the Docker Image:

    Build the Docker image using the provided Dockerfile:

    sudo docker build -t llama-embeddings .

  3. Run the Docker Container:

    Run the Docker container, mapping the container's port 8089 to the host's port 8089:

    sudo docker run -p 8089:8089 llama-embeddings

  4. Accessing the Application:

    The FastAPI application will now be accessible at http://localhost:8089 or at the static IP address of your VPS instance if you're running on one (You can get a 10-core, 30gb RAM, 1tb SSD with a static IP running Ubuntu 22.04 at Contabo for around $30/month, which is the cheapest I've found so far).

    You can interact then with the API using tools like curl or by accessing the FastAPI documentation at http://localhost:8089/docs.

  5. Viewing Logs:

    Logs from the application can be viewed directly in the terminal where you ran the docker run command.

Stopping and Managing the Container

Startup Procedures

During startup, the application performs the following tasks:

  1. Database Initialization:
    • The application initializes the SQLite database, setting up tables and executing important PRAGMAs to optimize performance.
    • Some of the important SQLite PRAGMAs include setting the database to use Write-Ahead Logging (WAL) mode, setting synchronous mode to NORMAL, increasing cache size to 1GB, setting the busy timeout to 2 seconds, and setting the WAL autocheckpoint to 100.
  2. Initialize Database Writer:
    • A dedicated database writer (DatabaseWriter) is initialized with a dedicated asynchronous queue to handle the write operations.
    • A set of hashes is created which represents the operations that are currently being processed or have already been processed. This avoids any duplicate operations in the queue.
  3. RAM Disk Setup:
    • If the USE_RAMDISK variable is enabled and the user has the required permissions, the application sets up a RAM Disk.
    • The application checks if there's already a RAM Disk set up at the specified path, if not, it calculates the optimal size for the RAM Disk and sets it up.
    • If the RAM Disk is enabled but the user lacks the required permissions, the RAM Disk feature is disabled and the application proceeds without it.
  4. Model Downloads:
    • The application downloads the required models.
  5. Model Loading:
    • Each downloaded model is loaded into memory. If any model file is not found, an error log is recorded.
  6. Build FAISS Indexes:
    • The application creates FAISS indexes for efficient similarity search using the embeddings from the database.
    • Associated texts are stored by model name for further use.

Note:

Endpoint Functionality and Workflow Overview

Here's a detailed breakdown of the main endpoints provided by the FastAPI server, explaining their functionality, input parameters, and how they interact with underlying models and systems:

1. /get_embedding_vector_for_string/ (POST)

Purpose

Retrieve the embedding vector for a given input text string using the specified model.

Parameters

Workflow

  1. Retrieve Embedding: The function retrieves or computes the embedding vector for the provided text using the specified or default model.
  2. Return Result: The embedding vector for the input text string is returned in the response.

2. /compute_similarity_between_strings/ (POST)

Purpose

Compute the similarity between two given input strings using specified model embeddings and a selected similarity measure.

Parameters

Workflow

  1. Retrieve Embeddings: The embeddings for text1 and text2 are retrieved or computed using the specified or default model.
  2. Compute Similarity: The similarity between the two embeddings is calculated using the specified similarity measure.
  3. Return Result: The similarity score, along with the embeddings and input texts, is returned in the response.

3. /search_stored_embeddings_with_query_string_for_semantic_similarity/ (POST)

Purpose

Find the most similar strings in the database to the given input "query" text. This endpoint uses a pre-computed FAISS index to quickly search for the closest matching strings.

Parameters

Workflow

  1. Search FAISS Index: The FAISS index, built on stored embeddings, is searched to find the most similar embeddings to the query_text.
  2. Return Result: The most similar strings found in the database, along with the similarity scores, are returned in the response.

4. /advanced_search_stored_embeddings_with_query_string_for_semantic_similarity/ (POST)

Purpose

Performs a two-step advanced semantic search. Utilizes FAISS and cosine similarity for initial filtering, followed by additional similarity measures for refined comparisons.

Parameters

Workflow

  1. Initial Filtering: Use FAISS and cosine similarity to find a set of similar strings.
  2. Refined Comparison: Apply additional similarity measures to the filtered set.
  3. Return Result: Return the most similar strings along with their multiple similarity scores.

Example Request

{
  "query_text": "Find me the most similar string!",
  "llm_model_name": "openchat_v3.2_super",
  "similarity_filter_percentage": 0.02,
  "number_of_most_similar_strings_to_return": 5
}

5. /get_all_embedding_vectors_for_document/ (POST)

Purpose

Extract text embeddings for a document. The library now supports a wide range of file types including plain text, .doc/.docx, PDF files, images (using Tesseract OCR), and many other types supported by the textract library.

Parameters

6. /compute_transcript_with_whisper_from_audio/ (POST)

Purpose

Transcribe an audio file and optionally compute document embeddings for the resulting transcript. This endpoint uses the Whisper model for transcription and a language model for generating embeddings. The transcription and embeddings can then be stored, and a ZIP file containing the embeddings can be made available for download.

Parameters

Request File and Parameters

You will need to use a multipart/form-data request to upload the audio file. The additional parameters like compute_embeddings_for_resulting_transcript_document and llm_model_name can be sent along as form fields.

Example Request

curl -X 'POST' \
  'http://localhost:8000/compute_transcript_with_whisper_from_audio/' \
  -H 'accept: application/json' \
  -H 'Authorization: Bearer YOUR_ACCESS_TOKEN' \
  -F 'file=@your_audio_file.wav' \
  -F 'compute_embeddings_for_resulting_transcript_document=true' \
  -F 'llm_model_name=custom-llm-model'

7. /get_text_completions_from_input_prompt/ (POST)

Purpose

Generate text completions for a given input prompt using the specified model.

Parameters

Request JSON Format

The JSON object should have the following keys:

Example Request

{
  "input_prompt": "The Kings of France in the 17th Century:",
  "llm_model_name": "phind-codellama-34b-python-v1",
  "temperature": 0.95,
  "grammar_file_string": "json",
  "number_of_tokens_to_generate": 500,
  "number_of_completions_to_generate": 3
}

8. /get_list_of_available_model_names/ (GET)

Purpose

Retrieve the list of available model names for generating embeddings.

Parameters

9. /get_all_stored_strings/ (GET)

Purpose

Retrieve a list of all stored strings from the database for which embeddings have been computed.

Parameters

10. /get_all_stored_documents/ (GET)

Purpose

Retrieve a list of all stored documents from the database for which embeddings have been computed.

Parameters

11. /clear_ramdisk/ (POST)

Purpose

Clear the RAM Disk to free up memory.

Parameters

12. /download/{file_name} (GET)

Purpose

Download a ZIP file containing document embeddings that were generated through the /compute_transcript_with_whisper_from_audio/ endpoint. The URL for this download will be supplied in the JSON response of the audio file transcription endpoint.

Parameters

13. /add_new_model/ (POST)

Purpose

Submit a new model URL for download and use. The model must be in .gguf format and larger than 100 MB to ensure it's a valid model file.

Parameters

Token-Level Embedding Vector Pooling

Pooling methods are designed to aggregate token-level embeddings, which are typically variable in length due to differing numbers of tokens in sentences or documents. By converting these token-level embeddings into a single, fixed-length vector, we ensure that each input text is represented consistently, regardless of its length. This fixed-length vector can then be used in various machine learning models that require inputs of a consistent size.

The primary goal of these pooling methods is to retain as much useful information as possible from the original token-level embeddings while ensuring that the transformation is deterministic and does not distort the data. Each method achieves this by applying different statistical or mathematical techniques to summarize the token embeddings.

Explanation of Pooling Methods

  1. SVD (Singular Value Decomposition):

    • How it works: Concatenates the first two singular vectors obtained from the SVD of the token embeddings matrix.
    • Rationale: SVD is a dimensionality reduction technique that captures the most important features of the data. Using the first two singular vectors provides a compact representation that retains significant information.

  2. SVD_First_Four:

    • How it works: Uses the first four singular vectors obtained from the SVD of the token embeddings matrix.
    • Rationale: By using more singular vectors, this method captures more of the variance in the data, providing a richer representation while still reducing dimensionality.

  3. ICA (Independent Component Analysis):

    • How it works: Applies ICA to the embeddings matrix to find statistically independent components, then flattens the result.
    • Rationale: ICA is useful for identifying independent sources in the data, providing a representation that highlights these independent features.

  4. Factor_Analysis:

    • How it works: Applies factor analysis to the embeddings matrix to identify underlying factors, then flattens the result.
    • Rationale: Factor analysis models the data in terms of latent factors, providing a summary that captures these underlying influences.

  5. Gaussian_Random_Projection:

    • How it works: Applies Gaussian random projection to reduce the dimensionality of the embeddings, then flattens the result.
    • Rationale: This method provides a fast and efficient way to reduce dimensionality while preserving the pairwise distances between points, useful for large datasets.