pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add Support for AWS Bedrock (#1408)

Browse Source 
### What problem does this PR solve?

#308 

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: KevinHuSh <kevinhu.sh@gmail.com>
This commit is contained in:
H
2024-07-08 09:37:34 +08:00
committed by GitHub
parent b3ebc66b13
commit 6144a109ab
8 changed files with 325 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
rag/llm/embedding_model.py
View File

@ -374,3 +374,48 @@ class MistralEmbed(Base):
res = self.client.embeddings(input=[truncate(text, 8196)],
model=self.model_name)
return np.array(res.data[0].embedding), res.usage.total_tokens
class BedrockEmbed(Base):
def __init__(self, key, model_name,
**kwargs):
import boto3
self.bedrock_ak = eval(key).get('bedrock_ak', '')
self.bedrock_sk = eval(key).get('bedrock_sk', '')
self.bedrock_region = eval(key).get('bedrock_region', '')
self.model_name = model_name
self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
def encode(self, texts: list, batch_size=32):
texts = [truncate(t, 8196) for t in texts]
embeddings = []
token_count = 0
for text in texts:
if self.model_name.split('.')[0] == 'amazon':
body = {"inputText": text}
elif self.model_name.split('.')[0] == 'cohere':
body = {"texts": [text], "input_type": 'search_document'}
response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body))
model_response = json.loads(response["body"].read())
embeddings.extend([model_response["embedding"]])
token_count += num_tokens_from_string(text)
return np.array(embeddings), token_count
def encode_queries(self, text):
embeddings = []
token_count = num_tokens_from_string(text)
if self.model_name.split('.')[0] == 'amazon':
body = {"inputText": truncate(text, 8196)}
elif self.model_name.split('.')[0] == 'cohere':
body = {"texts": [truncate(text, 8196)], "input_type": 'search_query'}
response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body))
model_response = json.loads(response["body"].read())
embeddings.extend([model_response["embedding"]])
return np.array(embeddings), token_count