SystemSh0cker/private-gpt
1
0
Fork 0
mirror of https://github.com/zylon-ai/private-gpt.git synced 2025-12-22 10:45:42 +01:00
Code Issues Projects Releases Packages Activity

Fix mypy

Browse source
This commit is contained in:
imartinez 2024-02-29 19:44:32 +01:00
parent a7b18058b5
commit 70ca241a8b
2 changed files with 6 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
private_gpt/components/llm/llm_component.py
View file

@ -113,7 +113,9 @@ class LLMComponent:
)
case "tensorrt":
try:
from llama_index.llms.nvidia_tensorrt import LocalTensorRTLLM # type: ignore
from llama_index.llms.nvidia_tensorrt import ( # type: ignore
LocalTensorRTLLM,
)
except ImportError as e:
raise ImportError(
"Nvidia TensorRTLLM dependencies not found, install with `poetry install --extras llms-nvidia-tensorrt`"

4
private_gpt/settings/settings.py
View file

@ -81,7 +81,9 @@ class DataSettings(BaseModel):
class LLMSettings(BaseModel):
mode: Literal["llamacpp", "openai", "openailike", "sagemaker", "mock", "ollama", "tensorrt"]
mode: Literal[
"llamacpp", "openai", "openailike", "sagemaker", "mock", "ollama", "tensorrt"
]
max_new_tokens: int = Field(
256,
description="The maximum number of token that the LLM is authorized to generate in one completion.",