Added api token for local llm
This commit is contained in:
@@ -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":
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user