copied encode function directly into search_embeddings.py

src/search_embeddings.py

@@ -2,13 +2,28 @@
 import numpy as np
 import polars as pl
 
-from src.encode import encode
 from sklearn.metrics.pairwise import cosine_similarity
 
 from transformers import AutoTokenizer, AutoModel
 
 import glob
 
+import torch
+
+#
+def encode(sentences, tokenizer, model, device="mps"):
+    inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt").to(device = device)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # outputs.last_hidden_state = [batch, tokens, hidden_dim]
+    # mean pooling
+    embeddings = outputs.last_hidden_state.mean(dim=1)
+
+    return(embeddings)
+
+
 
 # define the device where torch calculations take place
 my_device = "cpu"