diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2024-04-14 23:14:33 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2024-04-14 23:14:33 +0200 |
commit | db8b4891a0a46b2c6da31f1a50d6501e9904fd26 (patch) | |
tree | 508e0accfed6a09ab3dbdf78b5d9aa77addcf2cf /rag | |
parent | bd87b7fed75d5e504eb40c6616c2f1e1e56a0451 (diff) |
Format
Diffstat (limited to 'rag')
-rw-r--r-- | rag/generator/ollama.py | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/rag/generator/ollama.py b/rag/generator/ollama.py index b475dcf..77f8304 100644 --- a/rag/generator/ollama.py +++ b/rag/generator/ollama.py @@ -13,6 +13,7 @@ from .prompt import ANSWER_INSTRUCTION, Prompt class Ollama(metaclass=AbstractGenerator): def __init__(self) -> None: self.model = os.environ["GENERATOR_MODEL"] + log.debug(f"Using {self.model} for generator...") def __context(self, documents: List[Document]) -> str: results = [ @@ -39,7 +40,9 @@ class Ollama(metaclass=AbstractGenerator): for chunk in ollama.generate(model=self.model, prompt=metaprompt, stream=True): yield chunk["response"] - def chat(self, prompt: Prompt, messages: List[Dict[str, str]]) -> Generator[Any, Any, Any]: + def chat( + self, prompt: Prompt, messages: List[Dict[str, str]] + ) -> Generator[Any, Any, Any]: log.debug("Generating answer with ollama...") metaprompt = self.__metaprompt(prompt) messages.append({"role": "user", "content": metaprompt}) |