github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2025-08-25 18:49:23 +00:00
Code Issues Packages Projects Releases Wiki 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
Show Stats 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": case "tensorrt":
try: 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: except ImportError as e:
raise ImportError( raise ImportError(
"Nvidia TensorRTLLM dependencies not found, install with `poetry install --extras llms-nvidia-tensorrt`" "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): 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( max_new_tokens: int = Field(
256, 256,
description="The maximum number of token that the LLM is authorized to generate in one completion.", description="The maximum number of token that the LLM is authorized to generate in one completion.",