Merge pull request #46 from linto-ai/houssem_speaker_identify_optim

Use Qdrant in Speaker Identification

.envdefault

@@ -0,0 +1,9 @@
+SERVICE_MODE=http
+SERVICE_NAME=diarization 
+SERVICES_BROKER=redis://172.17.0.1:6379 
+BROKER_PASS=
+CONCURRENCY=2
+QDRANT_HOST=qdrant
+QDRANT_PORT=6333
+QDRANT_COLLECTION_NAME=speaker_embeddings
+QDRANT_RECREATE_COLLECTION=true

README.md

@@ -20,23 +20,101 @@ In what follow, you can replace "pyannote" by "simple" or "pybk" to try other me
 
 ### HTTP Server
 
-1. If needed, build docker image 
+1. If you want to use speaker identification, make sure Qdrant is running.
+First, create a custom bridge network so the diarization container can communicate with qdrant :
 
 ```bash
-docker build . -t linto-diarization-pyannote:latest -f pyannote/Dockerfile
+docker network create diarization_network
+```
+ You can start Qdrant using the following Docker command:
+
+```bash
+docker run 
+    --name qdrant \
+    --network diarization_network \
+    -p 6333:6333 \  # Qdrant default port
+    -v ./qdrant_storage:/qdrant/storage:z \
+    qdrant/qdrant
 ```
 
-2. Launch docker container (and keep it running)
+2. If needed, build docker image 
+
+```bash
+docker build . -t linto-diarization-pyannote:latest -f pyannote/Dockerfile
+```  
+
+3. Launch docker container (and keep it running)
+
+If you want to enable speaker identification, make sure to mount reference speaker audio samples to `/opt/speaker_samples`.
 
 ```bash
 docker run -it --rm \
+    --name linto-diarization \
+    --network diarization_network \
     -p 8080:80 \
+    -v ./data/speakers_samples:/opt/speaker_samples \ # Reference speaker samples. Enables speaker identification
     --shm-size=1gb --tmpfs /run/user/0 \
     --env SERVICE_MODE=http \
+    --env QDRANT_HOST=qdrant \ # Only specify if enabling speaker identification
+    --env QDRANT_PORT=6333 \ # Only specify if enabling speaker identification
+    --env QDRANT_COLLECTION_NAME=speaker_embeddings \ # Only specify if enabling speaker identification
+    --env QDRANT_RECREATE_COLLECTION=true \ # Only specify if enabling speaker identification
+    --env SERVICE_MODE=http \
     linto-diarization-pyannote:latest
 ```
 
-3. Open the swagger in a browser: [http://localhost:8080/docs](http://localhost:8080/docs)
+Alternatively, you can use docker-compose :
+
+```yaml
+
+services:
+  qdrant:
+    image: qdrant/qdrant
+    container_name: qdrant
+    ports:
+      - "6333:6333"  # Qdrant default port
+    volumes:
+      - ./qdrant_storage:/qdrant/storage:z
+
+  diarization_app: 
+    build: 
+      context : .
+      dockerfile: pyannote/Dockerfile
+    container_name: diarization_app
+    shm_size: '1gb'
+    stdin_open: true
+    tty: true     
+    ports :
+      - 8080:80
+    environment:
+      - QDRANT_HOST
+      - QDRANT_PORT
+      - QDRANT_COLLECTION_NAME
+      - QDRANT_RECREATE_COLLECTION
+      - SERVICE_MODE
+      - SERVICE_NAME
+      - SERVICES_BROKER
+      - CONCURRENCY
+    volumes:
+      - ./data/speakers_samples:/opt/speaker_samples # Reference Speaker samples : This enables speaker identification
+    depends_on:
+      - qdrant  # Ensure Qdrant starts before the app
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+
+```
+
+Run it using this command :
+```bash
+docker compose up
+```  
+
+4. Open the swagger in a browser: [http://localhost:8080/docs](http://localhost:8080/docs)
    Unfold `/diarization` route and click "Try it out". Then
    - Choose a file
    - Specify either `speaker_count` (Fixed number of speaker) or `max_speaker` (Max number of speakers)
@@ -52,7 +130,16 @@ In the following we assume we want to test on an audio that is in `$HOME/test.wa
 docker build . -t linto-diarization-pyannote:latest -f pyannote/Dockerfile
 ```
 
-2. Run Redis server
+2. If you want to use speaker identification, make sure Qdrant is running. You can start Qdrant using the following Docker command:
+
+```bash
+docker run 
+    -p 6333:6333 \  # Qdrant default port
+    -v ./qdrant_storage:/qdrant/storage:z \
+    qdrant/qdrant
+```
+
+3. Run Redis server
 
 ```bash
 docker run -it --rm \
@@ -61,7 +148,7 @@ docker run -it --rm \
     redis-server /etc/redis-stack.conf --protected-mode no --bind 0.0.0.0 --loglevel debug
 ```
 
-3. Launch docker container, attaching the volume where is the audio file on which you will test
+4. Launch docker container, attaching the volume where is the audio file on which you will test
 
 ```bash
 docker run -it --rm \
@@ -71,10 +158,14 @@ docker run -it --rm \
     --env SERVICES_BROKER=redis://172.17.0.1:6379 \
     --env BROKER_PASS= \
     --env CONCURRENCY=2 \
+    --env QDRANT_HOST=localhost \
+    --env QDRANT_PORT=6333 \
+    --env QDRANT_COLLECTION_NAME=speaker_embeddings \
+    --env QDRANT_RECREATE_COLLECTION=true \
     linto-diarization-pyannote:latest
 ```
 
-3. Testing with a given audio file can be done using python3 (with packages `celery` and `redis` installed).
+5. Testing with a given audio file can be done using python3 (with packages `celery` and `redis` installed).
    For example with the following command for the file `$HOME/test.wav` with 2 speakers
 
 ```bash

docker-compose.yml

@@ -1,16 +1,40 @@
-version: '3.7'
-
 services:
-  my-diarization-service:
-    image: linto-diarization-simple:latest
+  qdrant:
+    image: qdrant/qdrant
+    container_name: qdrant
+    ports:
+      - "6333:6333"  # Qdrant default port
     volumes:
-      - /path/to/shared/folder:/opt/audio
-    env_file: .env
-    deploy:
-      replicas: 1
-    networks:
-      - your-net
+      - ./qdrant_storage:/qdrant/storage:z
 
-networks:
-  your-net:
-    external: true
+  diarization_app: 
+    build: 
+      context : .
+      dockerfile: pyannote/Dockerfile
+    container_name: diarization_app
+    shm_size: '1gb'
+    stdin_open: true
+    tty: true     
+    ports :
+      - 8080:80
+    environment:
+      - QDRANT_HOST
+      - QDRANT_PORT
+      - QDRANT_COLLECTION_NAME
+      - QDRANT_RECREATE_COLLECTION
+      - SERVICE_MODE
+      - SERVICE_NAME
+      - SERVICES_BROKER
+      - CONCURRENCY
+    volumes:
+      - ./data/speakers_samples:/opt/speaker_samples # Reference Speaker samples 
+      - ./data/test_samples:/opt/audio # Test audio file (Celery task mode)
+    depends_on:
+      - qdrant  # Ensure Qdrant starts before the app
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

docker-entrypoint.sh

@@ -22,6 +22,21 @@ check_gpu_availability() {
 
 }
 
+# Wait for Qdrant to be available
+wait_for_qdrant() {
+    # Check if QDRANT_HOST and QDRANT_PORT are set
+    if [[ -z "${QDRANT_HOST}" || -z "${QDRANT_PORT}" ]]; then
+        echo "Qdrant environment variables are not set. Skipping wait for Qdrant."
+        return 0
+    fi
+    echo "Waiting for Qdrant to be reachable..."
+    /usr/src/app/wait-for-it.sh "${QDRANT_HOST}:${QDRANT_PORT}" --timeout=20 --strict -- echo "Qdrant is up"
+    if [ $? -ne 0 ]; then
+        echo "ERROR: Qdrant service not reachable at ${QDRANT_HOST}:${QDRANT_PORT}"
+        exit 1
+    fi
+}
+
 run_http_server() {
     echo "HTTP server Mode"
     python http_server/ingress.py --debug
@@ -58,6 +73,7 @@ run_celery_worker() {
 
 # Main logic
 check_gpu_availability
+wait_for_qdrant
 
 if [ -z "$SERVICE_MODE" ]; then
     echo "ERROR: Must specify a serving mode: [ http | task ]"