ALAimTrain / app.py
sungyi654's picture
Update app.py
4ba990c verified
# Import libraries.
import gradio as gr
import os
import pprint
import requests
import sys
import weaviate
from openai import OpenAI
from gradio.themes.base import Base
from icecream import ic
from langchain.callbacks.tracers import ConsoleCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFacePipeline
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.document_loaders import TextLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Weaviate
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from weaviate.embedded import EmbeddedOptions
def AIAimAdvice(query):
# OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
directory_path = "ApexDocs"
my_txts = os.listdir(directory_path)
# Load the TXT.
loaders = []
for my_txt in my_txts:
my_txt_path = os.path.join(directory_path, my_txt)
loaders.append(TextLoader(my_txt_path))
print("len(loaders) =", len(loaders))
# Load the TXT.
# data = [loader.load() for loader in loaders]
data = []
for loader in loaders:
data.append(loader.load())
print("len(data) =", len(data), "\n")
# Initialize the text splitter
# Uses a text splitter to split the data into smaller documents.
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
# docs = [text_splitter.split_documents(doc) for doc in data]
docs = []
for doc in data:
chunk = text_splitter.split_documents(doc)
docs.append(chunk)
# Debugging purposes.
# Print the number of total documents to be stored in the vector database.
total = 0
for i in range(len(docs)):
if i == len(docs) - 1:
print(len(docs[i]), end="")
else:
print(len(docs[i]), "+ " ,end="")
total += len(docs[i])
print(" =", total, " total documents\n")
# Print the first document.
print(docs[0], "\n\n\n")
# Print the total number of PDF files.
# docs is a list of lists where each list stores all the documents for one PDF file.
print(len(docs))
# Merge the documents to be embededed and store them in the vector database.
merged_documents = []
for doc in docs:
merged_documents.extend(doc)
# Print the merged list of all the documents.
print("len(merged_documents) =", len(merged_documents))
print(merged_documents)
# Hugging Face model for embeddings.
model_name = "sentence-transformers/all-MiniLM-L6-v2"
model_kwargs = {'device': 'cpu'}
# model_kwargs = {'device': 'cuda'}
embeddings = HuggingFaceEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
)
# Create Weaviate vector store (database).
client = weaviate.Client(
embedded_options = EmbeddedOptions()
)
# Initialize the Weaviate vector search with the document segments.
# Create a vector store (database) named vector_search from the sample documents.
vector_search = Weaviate.from_documents(
client = client,
documents = merged_documents,
# embedding = OpenAIEmbeddings(),
embedding = embeddings,
by_text = False
)
# Instantiate Weaviate Vector Search as a retriever
# Basic RAG.
# k to search for only the 10 most relevant documents.
# score_threshold to use only documents with a relevance score above 0.89.
retriever = vector_search.as_retriever(
search_type = "similarity", # similarity, mmr, similarity_score_threshold. https://api.python.langchain.com/en/latest/vectorstores/langchain_core.vectorstores.VectorStore.html#langchain_core.vectorstores.VectorStore.as_retriever
search_kwargs = {"k": 10, "score_threshold": 0.89}
)
# Define a prompt template.
# Define a LangChain prompt template to instruct the LLM to use our documents as the context.
# LangChain passes these documents to the {context} input variable and the user's query to the {question} variable.
template = """
You are an aim tutor for providing advice on improving aim specifically on Apex Legends.
Use the following pieces of retrieved context to answer the question at the end.
If you don't know the answer, just say that you don't know.
Context: {context}
Question: {question}
"""
custom_rag_prompt = ChatPromptTemplate.from_template(template)
# custom_rag_prompt = ChatPromptTemplate(
# template=template, input_variables=["context", "question"]
# )
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.2) # Increasing the temperature, the model becomes more creative and takes longer for inference.
# Construct a chain to answer questions on our data:
# 1. Weaviate Vector Search as the retriever to search for documents that are used as context by the LLM.
# 2. The prompt template that we constructed.
# 3. OpenAI's/Hugging Face chat model as the LLM used to generate a context-aware response.
# Regular chain format: chain = prompt | model | output_parser
rag_chain = (
{"context": retriever, "question": RunnablePassthrough()}
| custom_rag_prompt
| llm
# | hf
| StrOutputParser()
)
# First TXT file.
return rag_chain.invoke(query)
ApexAiming = gr.Interface(fn=AIAimAdvice, inputs=["text"], outputs="text", title="Apex Aiming advice")
ApexAiming.launch()