Spaces:
Sleeping
Sleeping
File size: 2,320 Bytes
eb957df |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex
import os
from llama_index.llms.openai import OpenAI
from llama_index.core import StorageContext, Settings, load_index_from_storage
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from retrieve import get_latest_dir
from datetime import datetime
load_dotenv()
Settings.llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
Settings.chunk_size = 2048
Settings.chunk_overlap = 24
def create_embedding():
"""
Create an embedding from the given directory.
Returns:
VectorStoreIndex: The index of the embedding from docs in the directory.
"""
output_dir = os.getenv("EMBEDDING_DIR")
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
embedding_path = f"{output_dir}/{timestamp}"
documents = SimpleDirectoryReader(os.getenv("PROD_SPEC_DIR")).load_data()
index = VectorStoreIndex.from_documents(documents, show_progress=True)
index.storage_context.persist(persist_dir=embedding_path)
return index
def load_embedding():
"""
Load the latest embedding from the directory.
Returns:
VectorStoreIndex: The index of the embedding from the latest directory.
"""
PERSIST_DIR = get_latest_dir(os.getenv("EMBEDDING_DIR"))
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)
return index
def query_rag_qa(rag_index, query, search_level):
"""
Query the RAG model for a given query.
Args:
rag_index (VectorStoreIndex): The RAG model index.
query (str): The query to ask the RAG model.
search_level (int): The max search level to use for the RAG model.
Returns:
tuple: The response, nodes, and reference text from the RAG model.
"""
myretriever = rag_index.as_retriever(
include_text=True,
similarity_top_k=search_level,
)
query_engine = rag_index.as_query_engine(
sub_retrievers=[
myretriever,
],
include_text=True,
similarity_top_k=search_level,
)
response = query_engine.query(query)
nodes = myretriever.retrieve(query)
reference_text = []
for node in nodes:
reference_text.append(node.text)
return response, nodes, reference_text
|