From 633f180eb0ccdc4772d5d705873cef1e33507976 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Fri, 5 Apr 2024 00:27:27 +0200 Subject: Wip llms --- rag/llms/__init__.py | 0 rag/llms/encoder.py | 24 ++++++++++++++++++++++++ rag/llms/generator.py | 0 3 files changed, 24 insertions(+) create mode 100644 rag/llms/__init__.py create mode 100644 rag/llms/encoder.py create mode 100644 rag/llms/generator.py (limited to 'rag/llms') diff --git a/rag/llms/__init__.py b/rag/llms/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rag/llms/encoder.py b/rag/llms/encoder.py new file mode 100644 index 0000000..758b523 --- /dev/null +++ b/rag/llms/encoder.py @@ -0,0 +1,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) + diff --git a/rag/llms/generator.py b/rag/llms/generator.py new file mode 100644 index 0000000..e69de29 -- cgit v1.2.3-70-g09d2