File size: 1,717 Bytes
220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 220c3c1 712f406 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | # query.py
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
# Set your Hugging Face API token
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
# Path to the Chroma database
CHROMA_PATH = "chroma"
# Prompt template for the LLM
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):
# Initialize embedding function and vector store
embedding_function = get_embedding()
db = Chroma(
persist_directory=CHROMA_PATH,
embedding_function=embedding_function
)
# Retrieve relevant documents
docs = db.similarity_search(query_text, k=3)
context = "\n\n".join([doc.page_content for doc in docs])
# Prepare prompt
prompt = PromptTemplate(
input_variables=["context", "question"],
template=PROMPT_TEMPLATE
).format(context=context, question=query_text)
# Initialize LLM
llm = HuggingFaceHub(
repo_id="google/flan-t5-base",
model_kwargs={"temperature": 0.5, "max_length": 512},
huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN
)
# Generate response
response = llm(prompt)
return response.strip()
if __name__ == "__main__":
main()
|