AskTRS / app.py
ucalyptus's picture
Update app.py
6e8a8b9
import whisper
import pytube
import gradio as gr
import openai
import faiss
from datetime import datetime
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.faiss import FAISS
from langchain.chains import RetrievalQAWithSourcesChain
from langchain import OpenAI
from langchain.vectorstores.base import VectorStoreRetriever
import os
def get_answer(api_key, video_link, question):
os.environ["OPENAI_API_KEY"] = api_key
video = pytube.YouTube(video_link)
audio = video.streams.get_audio_only()
fn = audio.download(output_path="tmp.mp3")
model = whisper.load_model("base")
transcription = model.transcribe(fn)
res = transcription['text']
def store_segments(text):
segment_size = 1000
segments = [{'text': text[i:i+segment_size], 'start': i} for i in range(0, len(text), segment_size)]
texts = []
start_times = []
for segment in segments:
text = segment['text']
start = segment['start']
start_datetime = datetime.fromtimestamp(start)
formatted_start_time = start_datetime.strftime('%H:%M:%S')
texts.append(text)
start_times.append(formatted_start_time)
return texts, start_times
texts, start_times = store_segments(res)
text_splitter = CharacterTextSplitter(chunk_size=1500, separator="\n")
docs = []
metadatas = []
for i, d in enumerate(texts):
splits = text_splitter.split_text(d)
docs.extend(splits)
metadatas.extend([{"source": start_times[i]}] * len(splits))
embeddings = OpenAIEmbeddings()
store = FAISS.from_texts(docs, embeddings, metadatas=metadatas)
faiss.write_index(store.index, "docs.index")
retri = VectorStoreRetriever(vectorstore=store)
chain = RetrievalQAWithSourcesChain.from_llm(llm=OpenAI(temperature=0), retriever=retri)
result = chain({"question": question})
return result['answer'], result['sources']
iface = gr.Interface(
fn=get_answer,
inputs=["text", "text", "text"],
outputs=["text", "text"],
)
iface.queue().launch()