blob: 95f3c6ace4193dda76feca662ca76361976de7d1 (
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
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
import os
from typing import Iterator, List
from uuid import uuid4
import ollama
from langchain_core.documents import Document
from loguru import logger as log
from qdrant_client.http.models import StrictFloat
try:
from rag.db.vector import Point
except ModuleNotFoundError:
from db.vector import Point
class Encoder:
def __init__(self) -> None:
self.model = os.environ["ENCODER_MODEL"]
self.query_prompt = "Represent this sentence for searching relevant passages: "
def __encode(self, prompt: str) -> List[StrictFloat]:
return list(ollama.embeddings(model=self.model, prompt=prompt)["embedding"])
def encode_document(self, chunks: Iterator[Document]) -> List[Point]:
log.debug("Encoding document...")
return [
Point(
id=uuid4().hex,
vector=self.__encode(chunk.page_content),
payload={"text": chunk.page_content},
)
for chunk in chunks
]
def encode_query(self, query: str) -> List[StrictFloat]:
log.debug(f"Encoding query: {query}")
query = self.query_prompt + query
return self.__encode(query)
|