From cc3bf37ca66d473a05d94aa879876b426908ed3e Mon Sep 17 00:00:00 2001 From: Robby Date: Sat, 4 May 2024 14:17:00 -0400 Subject: [PATCH] fmt --- llm_watsonx.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/llm_watsonx.py b/llm_watsonx.py index e71dd35..414605e 100644 --- a/llm_watsonx.py +++ b/llm_watsonx.py @@ -1,5 +1,5 @@ import os -from typing import List, Literal, Optional, Union, get_args, get_origin, Dict, Any +from typing import Any, Dict, List, Literal, Optional, Union, get_args, get_origin import llm from ibm_watsonx_ai.foundation_models import Embeddings, ModelInference, get_model_specs @@ -27,9 +27,11 @@ def get_env(): return (api_key, project_id) + def add_model_name_prefix(model): return watsonx_model_name_prefix + model + def strip_model_name_prefix(model): return model.lstrip(watsonx_model_name_prefix) @@ -182,7 +184,11 @@ def execute(self, prompt, stream, response, conversation): if prompt.system: prompt.prompt = prompt.system + "\n\n" + prompt.prompt - text = prompt.prompt if not conversation else self.build_chat_prompt(prompt, conversation) + text = ( + prompt.prompt + if not conversation + else self.build_chat_prompt(prompt, conversation) + ) params = prompt.options.to_payload()