| |
|
|
| import argparse |
| from langchain.vectorstores import Chroma |
| from langchain.prompts import PromptTemplate |
| from langchain.llms import HuggingFaceHub |
| from embedding_function import get_embedding |
| import os |
|
|
| |
| HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN") |
|
|
| |
| CHROMA_PATH = "chroma" |
|
|
| |
| PROMPT_TEMPLATE = """ |
| You are a helpful assistant. Use the following context to answer the question concisely and accurately. |
| |
| Context: |
| {context} |
| |
| Question: |
| {question} |
| |
| Answer: |
| """ |
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("query_text", type=str, help="The query text.") |
| args = parser.parse_args() |
| response = query_rag(args.query_text) |
| print("\n💬 Response:") |
| print(response) |
|
|
| def query_rag(query_text): |
| |
| embedding_function = get_embedding() |
| db = Chroma( |
| persist_directory=CHROMA_PATH, |
| embedding_function=embedding_function |
| ) |
|
|
| |
| docs = db.similarity_search(query_text, k=3) |
| context = "\n\n".join([doc.page_content for doc in docs]) |
|
|
| |
| prompt = PromptTemplate( |
| input_variables=["context", "question"], |
| template=PROMPT_TEMPLATE |
| ).format(context=context, question=query_text) |
|
|
| |
| llm = HuggingFaceHub( |
| repo_id="google/flan-t5-base", |
| model_kwargs={"temperature": 0.5, "max_length": 512}, |
| huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN |
| ) |
|
|
| |
| response = llm(prompt) |
| return response.strip() |
|
|
| if __name__ == "__main__": |
| main() |
|
|
