diff --git a/docker-compose-bedrock.yaml b/docker-compose-bedrock.yaml index 19ac5741..797fed2c 100644 --- a/docker-compose-bedrock.yaml +++ b/docker-compose-bedrock.yaml @@ -229,6 +229,12 @@ - "text-completion-bedrock" - "-p" - "pulsar://pulsar:6650" + - "-l" + - "4096" + - "-t" + - "0.0" + - "-m" + - "mistral.mixtral-8x7b-instruct-v0:1" - "-z" - "${AWS_ID_KEY}" - "-k" @@ -242,6 +248,12 @@ - "text-completion-bedrock" - "-p" - "pulsar://pulsar:6650" + - "-l" + - "4096" + - "-t" + - "0.0" + - "-m" + - "mistral.mixtral-8x7b-instruct-v0:1" - "-z" - "${AWS_ID_KEY}" - "-k" diff --git a/trustgraph/model/text_completion/azure/llm.py b/trustgraph/model/text_completion/azure/llm.py index 71338455..9ac20d84 100755 --- a/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph/model/text_completion/azure/llm.py @@ -19,6 +19,8 @@ module = ".".join(__name__.split(".")[1:-1]) default_input_queue = text_completion_request_queue default_output_queue = text_completion_response_queue default_subscriber = module +default_temperature = 0.0 +default_max = 4192 class Processor(ConsumerProducer): @@ -29,6 +31,8 @@ class Processor(ConsumerProducer): subscriber = params.get("subscriber", default_subscriber) endpoint = params.get("endpoint") token = params.get("token") + temperature = params.get("temperature", default_temperature) + max_tokens = params.get("max_output", default_max) super(Processor, self).__init__( **params | { @@ -54,8 +58,8 @@ class Processor(ConsumerProducer): "role": "user", "content": content } ], - "max_tokens": 4192, - "temperature": 0.2, + "max_tokens": max_tokens, + "temperature": temperature, "top_p": 1 } @@ -128,6 +132,18 @@ class Processor(ConsumerProducer): help=f'LLM model token' ) + parser.add_argument( + '-t', '--temperature', + default=f"temp=0.0", + help=f'LLM temperature parameter' + ) + + parser.add_argument( + '-l', '--max-output', + default=f"max_tokens=4192", + help=f'LLM max output tokens' + ) + def run(): Processor.start(module, __doc__) diff --git a/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph/model/text_completion/bedrock/llm.py index 11485af0..ea7c064d 100755 --- a/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph/model/text_completion/bedrock/llm.py @@ -22,6 +22,9 @@ default_output_queue = text_completion_response_queue default_subscriber = module default_model = 'mistral.mistral-large-2407-v1:0' default_region = 'us-west-2' +default_temperature = 0.0 +default_max = 2048 + class Processor(ConsumerProducer): @@ -34,6 +37,8 @@ class Processor(ConsumerProducer): aws_id = params.get("aws_id_key") aws_secret = params.get("aws_secret") aws_region = params.get("aws_region", default_region) + temperature = params.get("temperature", default_temperature) + max_tokens = params.get("max_output", default_max) super(Processor, self).__init__( **params | { @@ -74,8 +79,8 @@ class Processor(ConsumerProducer): if self.model.startswith("mistral"): promptbody = json.dumps({ "prompt": prompt, - "max_tokens": 8192, - "temperature": 0.0, + "max_tokens": max_tokens, + "temperature": temperature, "top_p": 0.99, "top_k": 40 }) @@ -84,8 +89,8 @@ class Processor(ConsumerProducer): elif self.model.startswith("meta"): promptbody = json.dumps({ "prompt": prompt, - "max_gen_len": 2048, - "temperature": 0.0, + "max_gen_len": max_tokens, + "temperature": temperature, "top_p": 0.95, }) @@ -93,8 +98,8 @@ class Processor(ConsumerProducer): elif self.model.startswith("anthropic"): promptbody = json.dumps({ "anthropic_version": "bedrock-2023-05-31", - "max_tokens": 8192, - "temperature": 0, + "max_tokens": max_tokens, + "temperature": temperature, "top_p": 0.999, "messages": [ { @@ -113,8 +118,8 @@ class Processor(ConsumerProducer): else: promptbody = json.dumps({ "prompt": prompt, - "max_tokens": 8192, - "temperature": 0.0, + "max_tokens": max_tokens, + "temperature": temperature, "top_p": 0.99, "top_k": 40 }) @@ -196,6 +201,18 @@ class Processor(ConsumerProducer): help=f'AWS Region (default: us-west-2)' ) + parser.add_argument( + '-t', '--temperature', + default=f"temp=0.0", + help=f'LLM temperature parameter' + ) + + parser.add_argument( + '-l', '--max-output', + default=f"max_tokens=2048", + help=f'LLM max output tokens' + ) + def run(): Processor.start(module, __doc__) diff --git a/trustgraph/model/text_completion/claude/llm.py b/trustgraph/model/text_completion/claude/llm.py index 29e5277d..e74db4c0 100755 --- a/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph/model/text_completion/claude/llm.py @@ -18,6 +18,8 @@ default_input_queue = text_completion_request_queue default_output_queue = text_completion_response_queue default_subscriber = module default_model = 'claude-3-5-sonnet-20240620' +default_temperature = 0.0 +default_max = 8192 class Processor(ConsumerProducer): @@ -28,6 +30,8 @@ class Processor(ConsumerProducer): subscriber = params.get("subscriber", default_subscriber) model = params.get("model", default_model) api_key = params.get("api_key") + temperature = params.get("temperature", default_temperature) + max_tokens = params.get("max_output", default_max) super(Processor, self).__init__( **params | { @@ -61,8 +65,8 @@ class Processor(ConsumerProducer): # FIXME: Rate limits? response = message = self.claude.messages.create( model=self.model, - max_tokens=1000, - temperature=0.1, + max_tokens=max_tokens, + temperature=temperature, system = "You are a helpful chatbot.", messages=[ { @@ -105,6 +109,18 @@ class Processor(ConsumerProducer): help=f'Claude API key' ) + parser.add_argument( + '-t', '--temperature', + default=f"temp=0.0", + help=f'LLM temperature parameter' + ) + + parser.add_argument( + '-l', '--max-output', + default=f"max_tokens=2048", + help=f'LLM max output tokens' + ) + def run(): Processor.start(module, __doc__)