llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added api token for local llm

Browse Source
This commit is contained in:
geens
2024-11-16 22:09:04 +01:00
parent 410e2b56f9
commit 506bc91639
2 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sia/__main__.py
View File

@@ -39,6 +39,7 @@ class Main:
self._llm = LocalLlmEngine( self._llm = LocalLlmEngine(
self._config.model, self._config.model,
self._config.temperature, self._config.temperature,
self._config.api_token,
self._config.token_limit self._config.token_limit
) )
case "hf": case "hf":

3
sia/local_llm_engine.py
View File

@@ -13,6 +13,7 @@ class LocalLlmEngine(LlmEngine):
model_path: str, model_path: str,
temperature: float, temperature: float,
token_limit: int, token_limit: int,
api_token: Optional[str],
): ):
""" """
Initialize the LLM Engine with a model path. Initialize the LLM Engine with a model path.
@@ -20,6 +21,7 @@ class LocalLlmEngine(LlmEngine):
Args: Args:
model_path: Path to the model weights to be used. model_path: Path to the model weights to be used.
temperature: Temperature for sampling temperature: Temperature for sampling
api_token: Huggingface API key
token_limit: Maximum number of tokens to generate token_limit: Maximum number of tokens to generate
""" """
self._temperature = temperature self._temperature = temperature
@@ -32,6 +34,7 @@ class LocalLlmEngine(LlmEngine):
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="auto", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
token=api_token,
) )
if self._tokenizer.pad_token_id is None: if self._tokenizer.pad_token_id is None:
self._tokenizer.pad_token_id = self._tokenizer.eos_token_id self._tokenizer.pad_token_id = self._tokenizer.eos_token_id