embedding

src/use_llm.py

@@ -31,17 +31,12 @@ def main_generate(prompt, model=DEFAULT_MODEL, system_prompt="You are a helpful
 
 
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-MODEL_ID = "thenlper/gte-large"  # embedding model
+from sentence_transformers import SentenceTransformer
 
-client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
+model = SentenceTransformer("thenlper/gte-large")  # downloaded from Hugging Face
 
 def get_embeddings(texts):
     if isinstance(texts, str):
         texts = [texts]
-    embeddings = []
-    for text in texts:
-        response = client.text_to_vector(text)
-        # response is usually a list of floats (the embedding vector)
-        embeddings.append(response)
+    embeddings = model.encode(texts, convert_to_numpy=True)
     return embeddings