Update codebase to support bedrock

weni-ai · Aug 29, 2024 · a908464 · a908464
1 parent 574223e
commit a908464
Show file tree

Hide file tree

Showing 4 changed files with 1,863 additions and 2,008 deletions.
diff --git a/bothub/api/v2/zeroshot/usecases/invoke_model.py b/bothub/api/v2/zeroshot/usecases/invoke_model.py
@@ -89,15 +89,15 @@ def _invoke_bedrock(self, prompt):
             "prompt": prompt
         })
 
-        response = bedrock.invoke_model(
+        bedrock_response = bedrock_runtime.invoke_model(
             body=payload,
             contentType='application/json',
             accept='application/json',
             modelId=settings.ZEROSHOT_BEDROCK_MODEL_ID,
             trace='ENABLED'
         )
 
-        classification = json.loads(response['body'].read().decode('utf-8'))
+        classification = json.loads(bedrock_response['body'].read().decode('utf-8'))
 
         classification_formatter = FormatClassification(classification)
         formatted_classification = classification_formatter.get_classification(self.zeroshot_data)

diff --git a/bothub/settings.py b/bothub/settings.py
@@ -130,7 +130,7 @@
     ZEROSHOT_TOP_K = (int, 10),
     ZEROSHOT_TEMPERATURE = (float, 0.1),
     ZEROSHOT_DO_SAMPLE = (bool, False),
-    ZEROSHOT_STOP = (str, "\n"),
+    ZEROSHOT_STOP = (list, ["<|end_of_text|>", "<|eot_id|>"]),
     ZEROSHOT_MODEL_BACKEND = (str, "runpod"),
 )