Spaces:
Running
Running
# Setup | |
# Import the necessary Libraries | |
import os | |
import json | |
import uuid | |
import gradio as gr | |
#import tiktoken | |
from datasets import load_dataset | |
#import pandas as pd | |
from openai import OpenAI | |
from langchain.text_splitter import RecursiveCharacterTextSplitter | |
from langchain_core.documents import Document | |
from langchain_community.embeddings.sentence_transformer import ( | |
SentenceTransformerEmbeddings | |
) | |
from langchain_community.vectorstores import Chroma | |
#from google.colab import userdata, drive | |
from langchain_community.document_loaders import PyPDFDirectoryLoader | |
#from google.colab import userdata | |
from huggingface_hub import CommitScheduler | |
from pathlib import Path | |
# Create Client | |
client = OpenAI( | |
base_url="https://api.openai.com/v1", | |
api_key=os.environ['CarlosGM'] | |
) | |
#api_key = os.environ.get("CarlosGM") | |
#client = OpenAI(api_key=api_key) | |
model_name = 'gpt-3.5-turbo' | |
# Define the embedding model and the vectorstore | |
embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large') | |
# Load the persisted vectorDB | |
## persisted_vectordb_location = '/content/drive/MyDrive/dataset_db' | |
dataset_10k_collection = 'Dataset-IBM-Meta-aws-google-msft' | |
vectorstore_persisted = Chroma( | |
collection_name=dataset_10k_collection, | |
persist_directory= './dataset_db', | |
embedding_function=embedding_model | |
) | |
retriever = vectorstore_persisted.as_retriever( | |
search_type='similarity', | |
search_kwargs={'k': 5} | |
) | |
# Prepare the logging functionality | |
log_file = Path("logs/") / f"data_{uuid.uuid4()}.json" | |
log_folder = log_file.parent | |
scheduler = CommitScheduler( | |
repo_id="10k-logs", | |
repo_type="dataset", | |
folder_path=log_folder, | |
path_in_repo="data", | |
every=2 | |
) | |
# Define the Q&A system message | |
qna_system_message = """ | |
You are an assistant to a financial services firm who answers user queries on annual 10 K reports. | |
User input will have the context required by you to answer user questions. | |
This context will begin with the token: ###Context. | |
The context contains references to specific portions of a document relevant to the user query. | |
The source for a context will begin with the token ###Source | |
User questions will begin with the token: ###Question. | |
Please answer only using the context provided in the input. Do not mention anything about the context in your final answer. | |
Please adhere to the following guidelines: | |
- Your response should only be about the question asked and nothing else. | |
- Answer only using the context provided. | |
- Do not mention anything about the context in your final answer. | |
- If the answer is not found in the context, it is very very important for you to respond with "I don't know. Please check the docs @ 'Dataset-10k file'" | |
- Always quote the source when you use the context. Cite the relevant source at the end of your response under the section - Source: | |
- Do not make up sources. Use the files provided in the sources section of the context and nothing else. You are prohibited from providing other sources. | |
If the answer is not found in the context, respond "I don't know". | |
Here is an example of how to structure your response: | |
Answer: | |
[Answer] | |
Source: | |
[Source] | |
""" | |
# Define the user message template | |
qna_user_message_template = """ | |
###Context | |
Here are some documents that are relevant to the question. | |
{context} | |
###Question | |
{question} | |
""" | |
# Define the predict function that runs when 'Submit' is clicked or when a API request is made | |
def predict(user_input,company): | |
#filter = "dataset/"+company+"-10-k-2023.pdf" | |
#relevant_document_chunks = vectorstore_persisted.similarity_search(user_input, k=5, filter={"source":filter}) | |
# Create context_for_query | |
filter = "dataset/"+company+"-10-k-2023.pdf" | |
relevant_document_chunks = retriever.invoke(user_input, k=5, filter={"source":filter}) | |
context_list = [d.page_content for d in relevant_document_chunks] | |
context_for_query = ". ".join(context_list) | |
# Create messages | |
prompt = [ | |
{'role':'system', 'content': qna_system_message}, | |
{'role': 'user', 'content': qna_user_message_template.format( | |
context=context_for_query, | |
question=user_input | |
) | |
} | |
] | |
# Get response from the LLM | |
try: | |
response = client.chat.completions.create( | |
model=model_name, | |
messages=prompt, | |
temperature=0 | |
) | |
prediction = response.choices[0].message.content.strip() | |
except Exception as e: | |
prediction = f'Sorry, I encountered the following error: \n {e}' | |
# While the prediction is made, log both the inputs and outputs to a local log file | |
# While writing to the log file, ensure that the commit scheduler is locked to avoid parallel | |
# access | |
with scheduler.lock: | |
with log_file.open("a") as f: | |
f.write(json.dumps( | |
{ | |
'user_input': [user_input, company], | |
'retrieved_context': context_for_query, | |
'model_response': prediction | |
} | |
)) | |
f.write("\n") | |
return prediction | |
# Set-up the Gradio UI | |
# Add text box and radio button to the interface | |
# The radio button is used to select the company 10k report in which the context needs to be retrieved. | |
textbox = gr.Textbox(placeholder='Enter your query here', lines=6) | |
company = gr.Radio(['aws', 'google', 'ibm', 'meta', 'msft'], label= "Select Company 10-k Report") | |
# Create the interface | |
demo = gr.Interface( | |
fn=predict, | |
inputs=[textbox,company], | |
outputs= "text", | |
title= "10-k Report Q&A", | |
description = "This Web API presents an inteface to ask questions about the 10-k reports" | |
) | |
# For the inputs parameter of Interface provide [textbox,company] | |
demo.queue() | |
demo.launch() | |