You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
37 lines
1.2 KiB
37 lines
1.2 KiB
from ollama import Client |
|
import os |
|
import env_manager |
|
|
|
env_manager.set_env() |
|
|
|
|
|
class LLM: |
|
def __init__( |
|
self, system_message=None, num_ctx=20000, temperature=0, chat=True |
|
) -> None: |
|
self.llm_model = "mistral-nemo:12b-instruct-2407-q5_K_M" #os.getenv("LLM_MODEL") |
|
self.system_message = system_message |
|
self.options = {"temperature": temperature, "num_ctx": num_ctx} |
|
self.messages = [{"role": "system", "content": self.system_message}] |
|
self.chat = chat |
|
self.ollama = Client( |
|
host=f'{os.getenv("LLM_URL")}:{os.getenv("LLM_PORT")}', |
|
) |
|
|
|
def generate(self, prompt: str) -> str: |
|
self.messages.append({"role": "user", "content": prompt}) |
|
|
|
result = self.ollama.chat( |
|
model=self.llm_model, messages=self.messages, options=self.options |
|
) |
|
|
|
answer = result["message"]["content"] |
|
self.messages.append({"role": "assistant", "content": answer}) |
|
if not self.chat: |
|
self.messages = [{"role": "system", "content": self.system_message}] |
|
|
|
return answer |
|
|
|
if __name__ == "__main__": |
|
llm = LLM() |
|
print(llm.generate("Why is the sky red?")) |