rag/llm/encoder.py


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

import os
from typing import List
from uuid import uuid4

import numpy as np
import ollama
from langchain_core.documents import Document
from qdrant_client.http.models import StrictFloat

from rag.db.embeddings 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: List[Document]) -> np.ndarray:
        return [
            Point(
                id=str(uuid4()),
                vector=self.__encode(chunk.page_content),
                payload={"text": chunk.page_content},
            )
            for chunk in chunks
        ]

    def query(self, query: str) -> np.ndarray:
        query = self.query_prompt + query
        return self.__encode(query)