|
| 1 | +from langchain.retrievers import AmazonKendraRetriever |
| 2 | +from langchain.chains import ConversationalRetrievalChain |
| 3 | +from langchain.prompts import PromptTemplate |
| 4 | +from langchain.llms import SagemakerEndpoint |
| 5 | +from langchain.llms.sagemaker_endpoint import LLMContentHandler |
| 6 | +import sys |
| 7 | +import json |
| 8 | +import os |
| 9 | + |
| 10 | +class bcolors: |
| 11 | + HEADER = '\033[95m' |
| 12 | + OKBLUE = '\033[94m' |
| 13 | + OKCYAN = '\033[96m' |
| 14 | + OKGREEN = '\033[92m' |
| 15 | + WARNING = '\033[93m' |
| 16 | + FAIL = '\033[91m' |
| 17 | + ENDC = '\033[0m' |
| 18 | + BOLD = '\033[1m' |
| 19 | + UNDERLINE = '\033[4m' |
| 20 | + |
| 21 | +MAX_HISTORY_LENGTH = 1 |
| 22 | + |
| 23 | +def build_chain(): |
| 24 | + region = os.environ["AWS_REGION"] |
| 25 | + kendra_index_id = os.environ["KENDRA_INDEX_ID"] |
| 26 | + endpoint_name = os.environ["LLAMA_2_ENDPOINT"] |
| 27 | + |
| 28 | + class ContentHandler(LLMContentHandler): |
| 29 | + content_type = "application/json" |
| 30 | + accepts = "application/json" |
| 31 | + |
| 32 | + def transform_input(self, prompt: str, model_kwargs: dict) -> bytes: |
| 33 | + input_str = json.dumps({"inputs": |
| 34 | + [[ |
| 35 | + #{"role": "system", "content": ""}, |
| 36 | + {"role": "user", "content": prompt}, |
| 37 | + ]], |
| 38 | + **model_kwargs |
| 39 | + }) |
| 40 | + #print(input_str) |
| 41 | + |
| 42 | + return input_str.encode('utf-8') |
| 43 | + |
| 44 | + def transform_output(self, output: bytes) -> str: |
| 45 | + response_json = json.loads(output.read().decode("utf-8")) |
| 46 | + |
| 47 | + return response_json[0]['generation']['content'] |
| 48 | + |
| 49 | + content_handler = ContentHandler() |
| 50 | + |
| 51 | + llm=SagemakerEndpoint( |
| 52 | + endpoint_name=endpoint_name, |
| 53 | + region_name=region, |
| 54 | + model_kwargs={"parameters": {"max_new_tokens": 1000, "top_p": 0.9,"temperature":0.6}}, |
| 55 | + endpoint_kwargs={"CustomAttributes":"accept_eula=true"}, |
| 56 | + content_handler=content_handler, |
| 57 | + ) |
| 58 | + |
| 59 | + retriever = AmazonKendraRetriever(index_id=kendra_index_id,region_name=region) |
| 60 | + |
| 61 | + prompt_template = """ |
| 62 | + Ignore everything before. |
| 63 | + |
| 64 | + Instruction: |
| 65 | + I want you to act as a research paper summarizer. I will provide you with a research paper on a specific topic in English, and you will create a summary. The summary should be concise and should accurately and objectively communicate the takeaway of the paper. You should not include any personal opinions or interpretations in your summary, but rather focus on objectively presenting the information from the paper. Your summary should be written in your own words and ensure that your summary is clear, concise, and accurately reflects the content of the original paper. |
| 66 | + |
| 67 | + First, provide a concise summary. Then provides the sources. |
| 68 | + |
| 69 | + {question} Answer "don't know" if not present in the document. |
| 70 | + {context} |
| 71 | + Solution:""" |
| 72 | + PROMPT = PromptTemplate( |
| 73 | + template=prompt_template, input_variables=["context", "question"], |
| 74 | + ) |
| 75 | + |
| 76 | + condense_qa_template = """ |
| 77 | + Chat History: |
| 78 | + {chat_history} |
| 79 | + Here is a new question for you: {question} |
| 80 | + Standalone question:""" |
| 81 | + standalone_question_prompt = PromptTemplate.from_template(condense_qa_template) |
| 82 | + |
| 83 | + qa = ConversationalRetrievalChain.from_llm( |
| 84 | + llm=llm, |
| 85 | + retriever=retriever, |
| 86 | + condense_question_prompt=standalone_question_prompt, |
| 87 | + return_source_documents=True, |
| 88 | + combine_docs_chain_kwargs={"prompt":PROMPT}, |
| 89 | + ) |
| 90 | + return qa |
| 91 | + |
| 92 | +def run_chain(chain, prompt: str, history=[]): |
| 93 | + print(prompt) |
| 94 | + return chain({"question": prompt, "chat_history": history}) |
| 95 | + |
| 96 | +if __name__ == "__main__": |
| 97 | + chat_history = [] |
| 98 | + qa = build_chain() |
| 99 | + print(bcolors.OKBLUE + "Hello! How can I help you?" + bcolors.ENDC) |
| 100 | + print(bcolors.OKCYAN + "Ask a question, start a New search: or CTRL-D to exit." + bcolors.ENDC) |
| 101 | + print(">", end=" ", flush=True) |
| 102 | + for query in sys.stdin: |
| 103 | + if (query.strip().lower().startswith("new search:")): |
| 104 | + query = query.strip().lower().replace("new search:","") |
| 105 | + chat_history = [] |
| 106 | + elif (len(chat_history) == MAX_HISTORY_LENGTH): |
| 107 | + chat_history.pop(0) |
| 108 | + result = run_chain(qa, query, chat_history) |
| 109 | + chat_history.append((query, result["answer"])) |
| 110 | + print(bcolors.OKGREEN + result['answer'] + bcolors.ENDC) |
| 111 | + if 'source_documents' in result: |
| 112 | + print(bcolors.OKGREEN + 'Sources:') |
| 113 | + for d in result['source_documents']: |
| 114 | + print(d.metadata['source']) |
| 115 | + print(bcolors.ENDC) |
| 116 | + print(bcolors.OKCYAN + "Ask a question, start a New search: or CTRL-D to exit." + bcolors.ENDC) |
| 117 | + print(">", end=" ", flush=True) |
| 118 | + print(bcolors.OKBLUE + "Bye" + bcolors.ENDC) |
| 119 | + |
| 120 | + |
| 121 | + |
0 commit comments