Update sentence_embeddings.py

sentence_embeddings.py

@@ -7,6 +7,7 @@ from datetime import datetime
 from logger import log
 from config import TEST_MODE
 
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 router = APIRouter()
 
 class SentenceEmbeddingsInput(BaseModel):
@@ -59,11 +60,11 @@ def generic_sentence_embeddings(model_name: str):
             tokenizer, model = loaded_models[model_name]
         else:
             tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModel.from_pretrained(model_name)
+            model = AutoModel.from_pretrained(model_name).to(device)
             loaded_models[model] = (tokenizer, model)
 
         # Tokenize sentences
-        encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
+        encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt').to(device)
         with torch.no_grad():
             model_output = model(**encoded_input)
             sentence_embeddings = model_output[0][:, 0]