From 506bc916395225a589efec344ed9f3f2bcc778e2 Mon Sep 17 00:00:00 2001
From: geens <niels.geens@gmail.com>
Date: Sat, 16 Nov 2024 22:09:04 +0100
Subject: [PATCH] Added api token for local llm

---
 sia/__main__.py         | 1 +
 sia/local_llm_engine.py | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/sia/__main__.py b/sia/__main__.py
index 2a80273..b667b2b 100644
--- a/sia/__main__.py
+++ b/sia/__main__.py
@@ -39,6 +39,7 @@ class Main:
                 self._llm = LocalLlmEngine(
                     self._config.model,
                     self._config.temperature,
+                    self._config.api_token,
                     self._config.token_limit
                 )
             case "hf":
diff --git a/sia/local_llm_engine.py b/sia/local_llm_engine.py
index 3e4840d..8aa9d9c 100644
--- a/sia/local_llm_engine.py
+++ b/sia/local_llm_engine.py
@@ -13,6 +13,7 @@ class LocalLlmEngine(LlmEngine):
         model_path: str,
         temperature: float,
         token_limit: int,
+        api_token: Optional[str],
     ):
         """
         Initialize the LLM Engine with a model path.
@@ -20,6 +21,7 @@ class LocalLlmEngine(LlmEngine):
         Args:
             model_path: Path to the model weights to be used.
             temperature: Temperature for sampling
+            api_token: Huggingface API key
             token_limit: Maximum number of tokens to generate
         """
         self._temperature = temperature
@@ -32,6 +34,7 @@ class LocalLlmEngine(LlmEngine):
             torch_dtype=torch.bfloat16,
             device_map="auto",
             trust_remote_code=True,
+            token=api_token,
         )
         if self._tokenizer.pad_token_id is None:
             self._tokenizer.pad_token_id = self._tokenizer.eos_token_id