Spaces:
Running
Running
import utils | |
import os | |
import openai | |
from llama_index import SimpleDirectoryReader | |
from llama_index import Document | |
from llama_index import VectorStoreIndex | |
from llama_index import ServiceContext | |
from llama_index.llms import OpenAI | |
from llama_index.embeddings import HuggingFaceEmbedding | |
from trulens_eval import Tru | |
from utils import get_prebuilt_trulens_recorder | |
import time | |
openai.api_key = utils.get_openai_api_key() | |
def main(): | |
if not os.path.exists("./default.sqlite"): | |
start_time = time.time() | |
file_ls_str = ", ".join(os.listdir("./raw_documents")) | |
print(f"File list: {file_ls_str}") | |
print("") | |
documents = SimpleDirectoryReader( | |
input_files=["./raw_documents/HI_Knowledge_Base.pdf"] | |
).load_data() | |
document = Document(text="\n\n".join([doc.text for doc in documents])) | |
### gpt-4-1106-preview | |
### gpt-3.5-turbo-1106 / gpt-3.5-turbo | |
print("Initializing GPT 3.5 ..") | |
llm = OpenAI(model="gpt-3.5-turbo-1106", temperature=0.1) | |
print("Initializing bge-small-en-v1.5 embedding model ..") | |
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5") | |
print("Creating vector store ..") | |
print("time spent:", time.time() - start_time) | |
service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model) | |
index = VectorStoreIndex.from_documents([document], service_context=service_context) | |
query_engine = index.as_query_engine() | |
separator = "\n\n" | |
eval_questions = [] | |
with open('./raw_documents/eval_questions.txt', 'r') as file: | |
content = file.read() | |
for question in content.split(separator): | |
print(question) | |
print(separator) | |
eval_questions.append(question.strip()) | |
response = query_engine.query(eval_questions[0]) | |
print(str(response)) | |
tru = Tru() | |
# tru.reset_database() | |
tru_recorder = get_prebuilt_trulens_recorder(query_engine, | |
app_id="Direct Query Engine") | |
print("Sending each question to llm ..") | |
with tru_recorder as recording: | |
for question in eval_questions: | |
response = query_engine.query(question) | |
records, feedback = tru.get_records_and_feedback(app_ids=[]) | |
os.makedirs("./results", exist_ok=True) | |
records.to_csv("./results/records.csv", index=False) | |
print(tru.db.engine.url.render_as_string(hide_password=False)) | |
end_time = time.time() | |
time_spent_mins = (end_time - start_time) / 60 | |
with open("./results/time_cost.txt", "w") as fp: | |
fp.write(f"Takes {int(time_spent_mins)} mins to create llm evaluation.") | |
# tru.run_dashboard() | |
if __name__ == "__main__": | |
main() |