neural_search.yaml

$schema: ../../../../json_schemas/test_story.schema.yaml

description: Test the creation a neural search ingest pipeline.
prologues:
  - path: /_cluster/settings
    method: PUT
    request:
      payload:
        persistent:
          plugins:
            ml_commons:
              only_run_on_ml_node: false
epilogues:
  - path: /_ingest/pipeline/movies_pipeline
    method: DELETE
    status: [200, 404]
  - path: /movies
    method: DELETE
    status: [200, 404]
  - path: /_plugins/_ml/models/{model_id}/_undeploy
    method: POST
    parameters:
      model_id: ${get_completed_register_model_task.model_id}
    status: [200, 404]
  - path: /_plugins/_ml/models/{model_id}
    parameters:
      model_id: ${get_completed_register_model_task.model_id}
    method: DELETE
    status: [200, 404]
  - path: /_plugins/_ml/model_groups/{model_group_id}
    method: DELETE
    status: [200, 404]
    parameters:
      model_group_id: ${create_model_group.test_model_group_id}
version: '>= 2.11'
chapters:
  - synopsis: Create model group.
    id: create_model_group
    path: /_plugins/_ml/model_groups/_register
    method: POST
    request:
      payload:
        name: NLP_Group
        description: Model group for NLP models.
    response:
      status: 200
    output:
      test_model_group_id: payload.model_group_id
  - synopsis: Register model.
    id: register_model
    path: /_plugins/_ml/models/_register
    method: POST
    request:
      payload:
        name: huggingface/sentence-transformers/msmarco-distilbert-base-tas-b
        version: 1.0.1
        model_format: TORCH_SCRIPT
    response:
      status: 200
    output:
      task_id: payload.task_id
  - synopsis: Wait to get completed task.
    id: get_completed_register_model_task
    path: /_plugins/_ml/tasks/{task_id}
    method: GET
    parameters:
      task_id: ${register_model.task_id}
    response:
      status: 200
      payload:
        state: COMPLETED
    output:
      model_id: payload.model_id
    retry:
      count: 3
      wait: 10000
  - synopsis: Deploy a model.
    id: deploy_model
    path: /_plugins/_ml/models/{model_id}/_deploy
    method: POST
    parameters:
      model_id: ${get_completed_register_model_task.model_id}
    output:
      task_id: payload.task_id
    response:
      status: 200
  - synopsis: Wait to get completed task.
    id: get_completed_deploy_model_task
    path: /_plugins/_ml/tasks/{task_id}
    method: GET
    parameters:
      task_id: ${deploy_model.task_id}
    response:
      status: 200
      payload:
        state: COMPLETED
    output:
      model_id: payload.model_id
    retry:
      count: 3
      wait: 10000      
  - synopsis: Create ingest pipeline for text embedding.
    path: /_ingest/pipeline/{id}
    method: PUT
    parameters:
      id: movies_pipeline
    request:
      payload:
        description: Extracts text from field and embeds it.
        processors:
          - text_embedding:
              model_id: text-embedding-model
              field_map:
                text: passage_embedding
    response:
      status: 200
      payload:
        acknowledged: true
  - synopsis: Create an index using the pipeline.
    path: /{index}
    method: PUT
    parameters:
      index: movies
    request:
      payload:
        settings:
          index.knn: true
          default_pipeline: movies_pipeline
        mappings:
          properties:
            title:
              type: text
            year:
              type: integer
            passage_embedding:
              type: knn_vector
              dimension: 768
              method:
                engine: lucene
                space_type: l2
                name: hnsw
                parameters: {}
    response:
      status: 200
      payload:
        acknowledged: true
  - synopsis: Ingest data.
    path: /_bulk
    method: POST
    parameters:
      refresh: 'true'
    request:
      content_type: application/x-ndjson
      payload:
        - {create: {_index: movies}}
        - {director: Bennett Miller, title: Moneyball, year: 2011}
        - {create: {_index: movies}}
        - {director: Nicolas Winding Refn, title: Drive, year: 1960}
    response:
      status: 200
  - synopsis: Search.
    path: /{index}/_search
    method: POST
    parameters:
      index: movies
    request:
      payload:
        _source:
          excludes: [passage_embedding]
        query:
          bool:
            should:
              - script_score:
                  query:
                    neural:
                      passage_embedding:
                        query_text: Money
                        model_id: ${get_completed_register_model_task.model_id}
                        k: 100
                  script:
                    source: _score * 1.5
              - script_score:
                  query:
                    match:
                      title: Moneyball
                  script:
                    source: _score * 1.7
  - synopsis: Undeploy a model.
    path: /_plugins/_ml/models/{model_id}/_undeploy
    method: POST
    parameters:
      model_id: ${get_completed_register_model_task.model_id}
    response:
      status: 200