Skip to content

Commit

Permalink
Merge pull request #40 from aws-samples/development
Browse files Browse the repository at this point in the history
Development
  • Loading branch information
MithilShah authored Sep 13, 2023
2 parents 409a3a9 + 3dbc203 commit 18ca7aa
Show file tree
Hide file tree
Showing 8 changed files with 392 additions and 32 deletions.
32 changes: 23 additions & 9 deletions kendra_retriever_samples/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,15 @@ If you are using Conda
conda env create -f environment.yml
```

### For Bedrock
If you are using Bedrock, make sure that you have a boto3 client with bedrock library and you use an AWS_PROFILE that has access to bedrock.

```
wget https://xxxxx/Documentation/SDK/bedrock-python-sdk.zip
unzip bedrock-python-sdk.zip
pip install *.whl
```

## Running samples
Before you run the sample, you need to deploy a Large Language Model (or get an API key if you using Anthropic or OPENAI). The samples in this repository have been tested on models deployed using SageMaker Jumpstart. The model id for the LLMS are specified in the table below.

Expand All @@ -40,6 +49,9 @@ Before you run the sample, you need to deploy a Large Language Model (or get an
| Flan XXL | FLAN_XXL_ENDPOINT | huggingface-text2text-flan-t5-xxl | flanxxl |
| Falcon 40B instruct | FALCON_40B_ENDPOINT | huggingface-llm-falcon-40b-instruct-bf16 | falcon40b |
| Llama2 70B instruct | LLAMA_2_ENDPOINT | meta-textgeneration-llama-2-70b-f | llama2 |
| Bedrock Titan | None | | bedrock_titan|
| Bedrock Claude | None | | bedrock_claude|
| Bedrock Claude V2 | None | | bedrock_claudev2|


after deploying the LLM, set up environment variables for kendra id, aws_region and the endpoint name (or the API key for an external provider)
Expand All @@ -49,15 +61,17 @@ For example, for running the `kendra_chat_flan_xl.py` sample, these environment
You can use commands as below to set the environment variables. Only set the environment variable for the provider that you are using. For example, if you are using Flan-xl only set the FLAN_XXL_ENDPOINT. There is no need to set the other Endpoints and keys.

```bash
export AWS_REGION="<YOUR-AWS-REGION>"
export KENDRA_INDEX_ID="<YOUR-KENDRA-INDEX-ID>"
export FLAN_XL_ENDPOINT="<YOUR-SAGEMAKER-ENDPOINT-FOR-FLAN-T-XL>" # only if you are using FLAN_XL
export FLAN_XXL_ENDPOINT="<YOUR-SAGEMAKER-ENDPOINT-FOR-FLAN-T-XXL>" # only if you are using FLAN_XXL
export FALCON_40B_ENDPOINT="<YOUR-SAGEMAKER-ENDPOINT-FOR-FALCON>" # only if you are using falcon as the endpoint
export LLAMA_2_ENDPOINT="<YOUR-SAGEMAKER-ENDPOINT-FOR-LLAMA2>" #only if you are using llama2 as the endpoint

export OPENAI_API_KEY="<YOUR-OPEN-AI-API-KEY>" # only if you are using OPENAI as the endpoint
export ANTHROPIC_API_KEY="<YOUR-ANTHROPIC-API-KEY>" # only if you are using Anthropic as the endpoint
export AWS_REGION=<YOUR-AWS-REGION>
export AWS_PROFILE=<AWS Profile>
export KENDRA_INDEX_ID=<YOUR-KENDRA-INDEX-ID>

export FLAN_XL_ENDPOINT=<YOUR-SAGEMAKER-ENDPOINT-FOR-FLAN-T-XL> # only if you are using FLAN_XL
export FLAN_XXL_ENDPOINT=<YOUR-SAGEMAKER-ENDPOINT-FOR-FLAN-T-XXL> # only if you are using FLAN_XXL
export FALCON_40B_ENDPOINT=<YOUR-SAGEMAKER-ENDPOINT-FOR-FALCON> # only if you are using falcon as the endpoint
export LLAMA_2_ENDPOINT=<YOUR-SAGEMAKER-ENDPOINT-FOR-LLAMA2> #only if you are using llama2 as the endpoint

export OPENAI_API_KEY=<YOUR-OPEN-AI-API-KEY> # only if you are using OPENAI as the endpoint
export ANTHROPIC_API_KEY=<YOUR-ANTHROPIC-API-KEY> # only if you are using Anthropic as the endpoint
```


Expand Down
27 changes: 26 additions & 1 deletion kendra_retriever_samples/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,11 @@
import kendra_chat_open_ai as openai
import kendra_chat_falcon_40b as falcon40b
import kendra_chat_llama_2 as llama2
import kendra_chat_bedrock_titan as bedrock_titan
import kendra_chat_bedrock_claude as bedrock_claude
import kendra_chat_bedrock_claudev2 as bedrock_claudev2



USER_ICON = "images/user-icon.png"
AI_ICON = "images/ai-icon.png"
Expand All @@ -21,6 +26,17 @@
'llama2' : 'Llama 2'
}

#function to read a properties file and create environment variables
def read_properties_file(filename):
import os
import re
with open(filename, 'r') as f:
for line in f:
m = re.match(r'^\s*(\w+)\s*=\s*(.*)\s*$', line)
if m:
os.environ[m.group(1)] = m.group(2)


# Check if the user ID is already stored in the session state
if 'user_id' in st.session_state:
user_id = st.session_state['user_id']
Expand Down Expand Up @@ -51,10 +67,19 @@
elif (sys.argv[1] == 'llama2'):
st.session_state['llm_app'] = llama2
st.session_state['llm_chain'] = llama2.build_chain()
elif (sys.argv[1] == 'bedrock_titan'):
st.session_state['llm_app'] = bedrock_titan
st.session_state['llm_chain'] = bedrock_titan.build_chain()
elif (sys.argv[1] == 'bedrock_claude'):
st.session_state['llm_app'] = bedrock_claude
st.session_state['llm_chain'] = bedrock_claude.build_chain()
elif (sys.argv[1] == 'bedrock_claudev2'):
st.session_state['llm_app'] = bedrock_claudev2
st.session_state['llm_chain'] = bedrock_claudev2.build_chain()
else:
raise Exception("Unsupported LLM: ", sys.argv[1])
else:
raise Exception("Usage: streamlit run app.py <anthropic|flanxl|flanxxl|openai>")
raise Exception("Usage: streamlit run app.py <anthropic|flanxl|flanxxl|openai|bedrock_titan|bedrock_claude|bedrock|claudev2>")

if 'chat_history' not in st.session_state:
st.session_state['chat_history'] = []
Expand Down
112 changes: 112 additions & 0 deletions kendra_retriever_samples/kendra_chat_bedrock_claude.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
# from aws_langchain.kendra import AmazonKendraRetriever #custom library
from langchain.retrievers import AmazonKendraRetriever
from langchain.chains import ConversationalRetrievalChain
from langchain.prompts import PromptTemplate
from langchain.llms.bedrock import Bedrock
from langchain.chains.llm import LLMChain
import sys
import os

class bcolors:
HEADER = '\033[95m'
OKBLUE = '\033[94m'
OKCYAN = '\033[96m'
OKGREEN = '\033[92m'
WARNING = '\033[93m'
FAIL = '\033[91m'
ENDC = '\033[0m'
BOLD = '\033[1m'
UNDERLINE = '\033[4m'

MAX_HISTORY_LENGTH = 5

def build_chain():
region = os.environ["AWS_REGION"]
kendra_index_id = os.environ["KENDRA_INDEX_ID"]
credentials_profile_name = os.environ['AWS_PROFILE']

print(credentials_profile_name)


llm = Bedrock(
credentials_profile_name=credentials_profile_name,
region_name = region,
model_kwargs={"max_tokens_to_sample":300,"temperature":1,"top_k":250,"top_p":0.999,"anthropic_version":"bedrock-2023-05-31"},
model_id="anthropic.claude-v1"
)

retriever = AmazonKendraRetriever(index_id=kendra_index_id,top_k=5,region_name=region)


prompt_template = """Human: This is a friendly conversation between a human and an AI.
The AI is talkative and provides specific details from its context but limits it to 240 tokens.
If the AI does not know the answer to a question, it truthfully says it
does not know.
Assistant: OK, got it, I'll be a talkative truthful AI assistant.
Human: Here are a few documents in <documents> tags:
<documents>
{context}
</documents>
Based on the above documents, provide a detailed answer for, {question}
Answer "don't know" if not present in the document.
Assistant:
"""
PROMPT = PromptTemplate(
template=prompt_template, input_variables=["context", "question"]
)

condense_qa_template = """Human:
Given the following conversation and a follow up question, rephrase the follow up question
to be a standalone question.
Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:
Assistant:"""
standalone_question_prompt = PromptTemplate.from_template(condense_qa_template)



qa = ConversationalRetrievalChain.from_llm(
llm=llm,
retriever=retriever,
condense_question_prompt=standalone_question_prompt,
return_source_documents=True,
combine_docs_chain_kwargs={"prompt":PROMPT},
verbose=True)

# qa = ConversationalRetrievalChain.from_llm(llm=llm, retriever=retriever, qa_prompt=PROMPT, return_source_documents=True)
return qa


def run_chain(chain, prompt: str, history=[]):
return chain({"question": prompt, "chat_history": history})


if __name__ == "__main__":
chat_history = []
qa = build_chain()
print(bcolors.OKBLUE + "Hello! How can I help you?" + bcolors.ENDC)
print(bcolors.OKCYAN + "Ask a question, start a New search: or CTRL-D to exit." + bcolors.ENDC)
print(">", end=" ", flush=True)
for query in sys.stdin:
if (query.strip().lower().startswith("new search:")):
query = query.strip().lower().replace("new search:","")
chat_history = []
elif (len(chat_history) == MAX_HISTORY_LENGTH):
chat_history.pop(0)
result = run_chain(qa, query, chat_history)
chat_history.append((query, result["answer"]))
print(bcolors.OKGREEN + result['answer'] + bcolors.ENDC)
if 'source_documents' in result:
print(bcolors.OKGREEN + 'Sources:')
for d in result['source_documents']:
print(d.metadata['source'])
print(bcolors.ENDC)
print(bcolors.OKCYAN + "Ask a question, start a New search: or CTRL-D to exit." + bcolors.ENDC)
print(">", end=" ", flush=True)
print(bcolors.OKBLUE + "Bye" + bcolors.ENDC)
112 changes: 112 additions & 0 deletions kendra_retriever_samples/kendra_chat_bedrock_claudev2.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
# from aws_langchain.kendra import AmazonKendraRetriever #custom library
from langchain.retrievers import AmazonKendraRetriever
from langchain.chains import ConversationalRetrievalChain
from langchain.prompts import PromptTemplate
from langchain.llms.bedrock import Bedrock
from langchain.chains.llm import LLMChain
import sys
import os

class bcolors:
HEADER = '\033[95m'
OKBLUE = '\033[94m'
OKCYAN = '\033[96m'
OKGREEN = '\033[92m'
WARNING = '\033[93m'
FAIL = '\033[91m'
ENDC = '\033[0m'
BOLD = '\033[1m'
UNDERLINE = '\033[4m'

MAX_HISTORY_LENGTH = 5

def build_chain():
region = os.environ["AWS_REGION"]
kendra_index_id = os.environ["KENDRA_INDEX_ID"]
credentials_profile_name = os.environ['AWS_PROFILE']

print(credentials_profile_name)


llm = Bedrock(
credentials_profile_name=credentials_profile_name,
region_name = region,
model_kwargs={"max_tokens_to_sample":300,"temperature":1,"top_k":250,"top_p":0.999,"anthropic_version":"bedrock-2023-05-31"},
model_id="anthropic.claude-v2"
)

retriever = AmazonKendraRetriever(index_id=kendra_index_id,top_k=5,region_name=region)


prompt_template = """Human: This is a friendly conversation between a human and an AI.
The AI is talkative and provides specific details from its context but limits it to 240 tokens.
If the AI does not know the answer to a question, it truthfully says it
does not know.
Assistant: OK, got it, I'll be a talkative truthful AI assistant.
Human: Here are a few documents in <documents> tags:
<documents>
{context}
</documents>
Based on the above documents, provide a detailed answer for, {question}
Answer "don't know" if not present in the document.
Assistant:
"""
PROMPT = PromptTemplate(
template=prompt_template, input_variables=["context", "question"]
)

condense_qa_template = """Human:
Given the following conversation and a follow up question, rephrase the follow up question
to be a standalone question.
Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:
Assistant:"""
standalone_question_prompt = PromptTemplate.from_template(condense_qa_template)



qa = ConversationalRetrievalChain.from_llm(
llm=llm,
retriever=retriever,
condense_question_prompt=standalone_question_prompt,
return_source_documents=True,
combine_docs_chain_kwargs={"prompt":PROMPT},
verbose=True)

# qa = ConversationalRetrievalChain.from_llm(llm=llm, retriever=retriever, qa_prompt=PROMPT, return_source_documents=True)
return qa


def run_chain(chain, prompt: str, history=[]):
return chain({"question": prompt, "chat_history": history})


if __name__ == "__main__":
chat_history = []
qa = build_chain()
print(bcolors.OKBLUE + "Hello! How can I help you?" + bcolors.ENDC)
print(bcolors.OKCYAN + "Ask a question, start a New search: or CTRL-D to exit." + bcolors.ENDC)
print(">", end=" ", flush=True)
for query in sys.stdin:
if (query.strip().lower().startswith("new search:")):
query = query.strip().lower().replace("new search:","")
chat_history = []
elif (len(chat_history) == MAX_HISTORY_LENGTH):
chat_history.pop(0)
result = run_chain(qa, query, chat_history)
chat_history.append((query, result["answer"]))
print(bcolors.OKGREEN + result['answer'] + bcolors.ENDC)
if 'source_documents' in result:
print(bcolors.OKGREEN + 'Sources:')
for d in result['source_documents']:
print(d.metadata['source'])
print(bcolors.ENDC)
print(bcolors.OKCYAN + "Ask a question, start a New search: or CTRL-D to exit." + bcolors.ENDC)
print(">", end=" ", flush=True)
print(bcolors.OKBLUE + "Bye" + bcolors.ENDC)
Loading

0 comments on commit 18ca7aa

Please sign in to comment.