Claudio Ortolina
97 lines
2.4 KiB
Elixir
97 lines
2.4 KiB
Elixir
defmodule Mix.Tasks.MusicLibrary.GenerateEmbeddings do
|
|
@moduledoc """
|
|
Generates embeddings for records.
|
|
|
|
## Usage
|
|
|
|
# Generate embeddings for all records without embeddings
|
|
mix music_library.generate_embeddings
|
|
|
|
# Regenerate embeddings for all records (force)
|
|
mix music_library.generate_embeddings --force
|
|
|
|
# Generate embeddings for specific record IDs
|
|
mix music_library.generate_embeddings --ids=id1,id2,id3
|
|
|
|
## Options
|
|
|
|
* `--force` - Regenerate embeddings for all records, even if they already exist
|
|
* `--ids` - Comma-separated list of record IDs to process
|
|
"""
|
|
use Mix.Task
|
|
|
|
import Ecto.Query
|
|
|
|
alias MusicLibrary.Repo
|
|
alias MusicLibrary.Records.{Record, RecordEmbedding}
|
|
alias MusicLibrary.Worker.GenerateRecordEmbedding
|
|
|
|
@shortdoc "Generates embeddings for records"
|
|
|
|
@impl Mix.Task
|
|
def run(args) do
|
|
Mix.Task.run("app.start")
|
|
|
|
{opts, _} =
|
|
OptionParser.parse!(args,
|
|
strict: [force: :boolean, ids: :string],
|
|
aliases: [f: :force, i: :ids]
|
|
)
|
|
|
|
force? = Keyword.get(opts, :force, false)
|
|
ids = parse_ids(Keyword.get(opts, :ids))
|
|
|
|
records = get_records(force?, ids)
|
|
total = length(records)
|
|
|
|
Mix.shell().info("Found #{total} records to process...")
|
|
|
|
records
|
|
|> Enum.with_index(1)
|
|
|> Enum.each(fn {record, index} ->
|
|
Mix.shell().info("[#{index}/#{total}] Enqueueing job for: #{record.title}")
|
|
|
|
%{record_id: record.id}
|
|
|> GenerateRecordEmbedding.new()
|
|
|> Oban.insert!()
|
|
end)
|
|
|
|
Mix.shell().info("\nSuccessfully enqueued #{total} embedding generation jobs.")
|
|
Mix.shell().info("Monitor progress in Oban dashboard or logs.")
|
|
end
|
|
|
|
defp parse_ids(nil), do: []
|
|
|
|
defp parse_ids(ids_string) do
|
|
ids_string
|
|
|> String.split(",")
|
|
|> Enum.map(&String.trim/1)
|
|
|> Enum.reject(&(&1 == ""))
|
|
end
|
|
|
|
defp get_records(force?, []) when force? do
|
|
Record
|
|
|> order_by(asc: :title)
|
|
|> Repo.all()
|
|
end
|
|
|
|
defp get_records(false, []) do
|
|
existing_record_ids =
|
|
RecordEmbedding
|
|
|> select([re], re.record_id)
|
|
|> Repo.all()
|
|
|> MapSet.new()
|
|
|
|
Record
|
|
|> order_by(asc: :title)
|
|
|> Repo.all()
|
|
|> Enum.reject(fn record -> MapSet.member?(existing_record_ids, record.id) end)
|
|
end
|
|
|
|
defp get_records(_force?, ids) when is_list(ids) and length(ids) > 0 do
|
|
Record
|
|
|> where([r], r.id in ^ids)
|
|
|> Repo.all()
|
|
end
|
|
end
|