blob: 758b52387c61a9411377e3953b3adad5e67cf3a2 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
|
from typing import List
import ollama
import numpy as np
# FIXME: .env
ENCODER_MODEL = "mxbai-embed-large"
class Encoder:
def __init__(self) -> None:
self.query_prompt = "Represent this sentence for searching relevant passages: "
def __encode(self, prompt: str) -> np.ndarray:
x = ollama.embeddings(model=ENCODER_MODEL, prompt=prompt)
x = np.array([x["embedding"]]).astype("float32")
return x
def encode(self, doc: List[str]) -> List[np.ndarray]:
return [self.__encode(chunk) for chunk in doc]
def query(self, query: str) -> np.ndarray:
query = self.query_prompt + query
return self.__encode(query)
|
