Added temperature argument

This commit is contained in:
Niels Geens
2024-11-13 11:51:02 +01:00
parent 2fa2eca9bf
commit aa37182614
2 changed files with 14 additions and 2 deletions

View File

@@ -48,7 +48,8 @@ class Main:
case "hf":
self._llm = HfLlmEngine(
model_id=self._config.model,
api_token=self._config.api_token
api_token=self._config.api_token,
temperature=self._config.temperature
)
case "test":
self._llm = TestLLM()

View File

@@ -75,6 +75,12 @@ class Config:
default=os.getenv('SIA_MODEL', '/root/model/'),
help='Path to the model directory (default: /root/model/, env: SIA_MODEL)'
)
parser.add_argument(
'--temperature',
type=float,
default=float(os.getenv('SIA_TEMPERATURE', '0.7')),
help='LLM temperature parameter (default: 0.7, env: SIA_TEMPERATURE)'
)
self.args = parser.parse_args()
def _parse_bool_env(self, env_var: str, default: bool) -> bool:
@@ -135,3 +141,8 @@ class Config:
def model(self) -> str:
"""Path to the model directory."""
return self.args.model
@property
def temperature(self) -> float:
"""LLM temperature parameter."""
return self.args.temperature