Ruby script using BERT to classify search queries as keywords, questions, or statements.

convert_to_onnx.sh

#!/usr/bin/env bash

uv tool install optimum[exporters]

model_name="bert-mini-finetune-question-detection"

optimum-cli export onnx --model "shahrukhx01/${model_name}" "${model_name}" --optimize O2

jq '. += {"id2label": {"0": "KEYWORD", "1": "QUESTION_OR_STATEMENT"}}' "${model_name}/config.json" > tmp.json && mv tmp.json "${model_name}/config.json"
jq '.model_max_length = 512' "${model_name}/tokenizer_config.json" > tmp.json && mv tmp.json "${model_name}/tokenizer_config.json"

Gemfile

gem "informers"

keyword_or_nl_query.rb

require 'informers'

classifier = Informers.pipeline(
  "text-classification", "bert-mini-finetune-question-detection", 
  cache_dir: File.expand_path(__dir__),
  local_files_only: true,
  model_file_name: "../model",
  quantized: false
)

queries = [
  # Keyword queries
  "ruby on rails installation",
  "postgresql database setup",
  "rspec test examples",
  
  # Question queries
  "How to install Ruby on Rails?",
  "What are the best practices for using Capybara?",
  "Can you explain how to integrate Elasticsearch with Rails?",
  
  # Statement queries
  "I need help with AWS SDK configuration in Rails.",
  "The Kafka integration is not working as expected.",
  "Webpack is throwing errors during asset compilation."
]

results = classifier.call(queries)

queries.zip(results) do |query, result|
  puts "Query: #{query}"
  puts "Label: #{result[:label]} (#{(result[:score] * 100).round(2)}% confidence)"
  puts
end

Usage

bin/bundle exec ruby keyword_or_nl_query.rb

Output:

Query: ruby on rails installation
Label: KEYWORD (82.05% confidence)

Query: postgresql database setup
Label: KEYWORD (99.92% confidence)

Query: rspec test examples
Label: KEYWORD (99.92% confidence)

Query: How to install Ruby on Rails?
Label: QUESTION_OR_STATEMENT (99.98% confidence)

Query: What are the best practices for using Capybara?
Label: QUESTION_OR_STATEMENT (99.98% confidence)

Query: Can you explain how to integrate Elasticsearch with Rails?
Label: QUESTION_OR_STATEMENT (99.98% confidence)

Query: I need help with AWS SDK configuration in Rails.
Label: QUESTION_OR_STATEMENT (99.93% confidence)

Query: The Kafka integration is not working as expected.
Label: QUESTION_OR_STATEMENT (99.98% confidence)

Query: Webpack is throwing errors during asset compilation.
Label: QUESTION_OR_STATEMENT (99.97% confidence)