|
import os |
|
import streamlit as st |
|
import pickle |
|
import time |
|
from langchain import OpenAI |
|
from langchain.chains import RetrievalQAWithSourcesChain |
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
from langchain.document_loaders import UnstructuredURLLoader |
|
from langchain.embeddings import FakeEmbeddings |
|
from langchain.llms import HuggingFaceHub |
|
from langchain.chains import LLMChain |
|
from langchain.vectorstores import FAISS |
|
|
|
from dotenv import load_dotenv |
|
load_dotenv() |
|
os.environ["HUGGINGFACEHUB_API_TOKEN"] = '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' |
|
|
|
|
|
class RockyBot: |
|
def __init__(self, llm): |
|
self.llm = llm |
|
self.vectorstore = None |
|
|
|
def process_urls(self, urls): |
|
"""Processes the given URLs and saves the FAISS index to a pickle file.""" |
|
|
|
|
|
loader = UnstructuredURLLoader(urls=urls) |
|
|
|
|
|
text_splitter = RecursiveCharacterTextSplitter( |
|
separators=['\n\n', '\n', '.', ','], |
|
chunk_size=1000 |
|
) |
|
docs = text_splitter.split_documents(loader.load()) |
|
|
|
|
|
embeddings = FakeEmbeddings(size=1352) |
|
self.vectorstore = FAISS.from_documents(docs, embeddings) |
|
|
|
|
|
with open("faiss_store_openai.pkl", "wb") as f: |
|
pickle.dump(self.vectorstore, f) |
|
|
|
def answer_question(self, question): |
|
"""Answers the given question using the LLM and retriever.""" |
|
|
|
chain = RetrievalQAWithSourcesChain.from_llm(llm=self.llm, retriever=self.vectorstore.as_retriever()) |
|
result = chain({"question": question}, return_only_outputs=True) |
|
|
|
return result["answer"], result.get("sources", "") |
|
|
|
|
|
if __name__ == '__main__': |
|
llm = HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature": 0.5, "max_length": 64}) |
|
rockybot = RockyBot(llm) |
|
|
|
|
|
if st.sidebar.button("Process URLs"): |
|
rockybot.process_urls(st.sidebar.text_input("URL 1"), st.sidebar.text_input("URL 2"), st.sidebar.text_input("URL 3")) |
|
st.progress(100.0) |
|
|
|
|
|
query = st.text_input("Question: ") |
|
if query: |
|
answer, sources = rockybot.answer_question(query) |
|
|
|
st.header("Answer") |
|
st.write(answer) |
|
|
|
|
|
if sources: |
|
st.subheader("Sources:") |
|
for source in sources.split("\n"): |
|
st.write(source) |
|
|
|
|