File size: 715 Bytes
9db894e |
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 |
from qa.chains import retrieval_qa
from qa.loader import youtube_doc_loader
from qa.model import load_llm
from qa.split import split_document
from qa.vector_store import create_vector_store
class YoutubeQA:
def __init__(self):
pass
def load_model(self) -> None:
self.llm = load_llm()
def load_vector_store(self, url: str) -> None:
data = youtube_doc_loader(url=url)
docs = split_document(data=data)
self.retriver = create_vector_store(docs=docs)
def load_retriever(self) -> None:
self.retrieval_qa = retrieval_qa(llm=self.llm, retriever=self.retriver)
def run(self, question: str) -> str:
return self.retrieval_qa.run(question)
|