# Load reticulate for Python integration
library(reticulate)

# Import necessary Python libraries
pd <- import("pandas")
torch <- import("torch")
sentence_transformers <- import("sentence_transformers")

# Load the sentence-transformer model
load_model <- function(model_name) {
  cat("Loading model:", model_name, "\n")
  embedder <- sentence_transformers$SentenceTransformer(model_name)
  
  # Check and use GPU if available
  if (torch$cuda$is_available()) {
    cat("GPU found! Using CUDA.\n")
    embedder$to("cuda")
  } else {
    cat("No GPU found. Using CPU.\n")
  }
  
  return(embedder)
}

# Function to generate embedding for a query
generate_query_embedding <- function(query, model_name = "sentence-transformers/all-MiniLM-L6-v2") {
  # Load the model
  embedder <- load_model(model_name)
  
  # Generate embedding
  cat("Generating embedding for the query...\n")
  embedding <- embedder$encode(query)
  
  # Convert the embedding to a numeric vector
  return(as.numeric(embedding))
}

# # Example usage
# query <- "I need a package for data visualization"
# model_name <- "sentence-transformers/all-MiniLM-L6-v2"
# query_embedding <- generate_query_embedding(query, model_name)
# 
# # Print the embedding
# cat("Embedding for query:\n")
# print(query_embedding)