From 1cf0a401054c3e3ebde60bfd73ad15e39bc531e6 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Fri, 5 Apr 2024 00:42:02 +0200 Subject: Rename llm --- rag/llm/encoder.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 rag/llm/encoder.py (limited to 'rag/llm/encoder.py') diff --git a/rag/llm/encoder.py b/rag/llm/encoder.py new file mode 100644 index 0000000..d5e0566 --- /dev/null +++ b/rag/llm/encoder.py @@ -0,0 +1,23 @@ +import os +from typing import List + +import numpy as np +import ollama + + +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) -> 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) -- cgit v1.2.3-70-g09d2