Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,84 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import whisper
|
2 |
+
import pytube
|
3 |
+
import gradio as gr
|
4 |
+
import openai
|
5 |
+
import faiss
|
6 |
+
from datetime import datetime
|
7 |
+
from langchain.embeddings.openai import OpenAIEmbeddings
|
8 |
+
from langchain.text_splitter import CharacterTextSplitter
|
9 |
+
from langchain.vectorstores.faiss import FAISS
|
10 |
+
from langchain.chains import RetrievalQAWithSourcesChain
|
11 |
+
from langchain import OpenAI
|
12 |
+
from langchain.vectorstores.base import VectorStoreRetriever
|
13 |
+
import os
|
14 |
+
|
15 |
+
video_data_cache = {}
|
16 |
+
|
17 |
+
def get_answer(api_key, video_link, question):
|
18 |
+
os.environ["OPENAI_API_KEY"] = api_key
|
19 |
+
|
20 |
+
if video_link not in video_data_cache:
|
21 |
+
video = pytube.YouTube(video_link)
|
22 |
+
audio = video.streams.get_audio_only()
|
23 |
+
fn = audio.download(output_path="tmp.mp3")
|
24 |
+
model = whisper.load_model("base")
|
25 |
+
transcription = model.transcribe(fn)
|
26 |
+
res = transcription['text']
|
27 |
+
|
28 |
+
def store_segments(text):
|
29 |
+
segment_size = 1000
|
30 |
+
segments = [{'text': text[i:i+segment_size], 'start': i} for i in range(0, len(text), segment_size)]
|
31 |
+
|
32 |
+
texts = []
|
33 |
+
start_times = []
|
34 |
+
|
35 |
+
for segment in segments:
|
36 |
+
text = segment['text']
|
37 |
+
start = segment['start']
|
38 |
+
|
39 |
+
start_datetime = datetime.fromtimestamp(start)
|
40 |
+
formatted_start_time = start_datetime.strftime('%H:%M:%S')
|
41 |
+
|
42 |
+
texts.append(text)
|
43 |
+
start_times.append(formatted_start_time)
|
44 |
+
|
45 |
+
return texts, start_times
|
46 |
+
|
47 |
+
texts, start_times = store_segments(res)
|
48 |
+
|
49 |
+
text_splitter = CharacterTextSplitter(chunk_size=1500, separator="\n")
|
50 |
+
docs = []
|
51 |
+
metadatas = []
|
52 |
+
for i, d in enumerate(texts):
|
53 |
+
splits = text_splitter.split_text(d)
|
54 |
+
docs.extend(splits)
|
55 |
+
metadatas.extend([{"source": start_times[i]}] * len(splits))
|
56 |
+
|
57 |
+
embeddings = OpenAIEmbeddings()
|
58 |
+
store = FAISS.from_texts(docs, embeddings, metadatas=metadatas)
|
59 |
+
faiss.write_index(store.index, f"docs.index")
|
60 |
+
|
61 |
+
video_data_cache[video_link] = f"docs.index"
|
62 |
+
|
63 |
+
index_file = video_data_cache[video_link]
|
64 |
+
store = faiss.read_index(index_file)
|
65 |
+
|
66 |
+
retri = VectorStoreRetriever(vectorstore=store)
|
67 |
+
|
68 |
+
chain = RetrievalQAWithSourcesChain.from_llm(llm=OpenAI(temperature=0), retriever=retri)
|
69 |
+
|
70 |
+
result = chain({"question": question})
|
71 |
+
|
72 |
+
return result['answer'], result['sources']
|
73 |
+
|
74 |
+
|
75 |
+
iface = gr.Interface(
|
76 |
+
fn=get_answer,
|
77 |
+
inputs=["text", "text", "text"],
|
78 |
+
outputs=["text", "text"],
|
79 |
+
examples=[
|
80 |
+
["sk-kVc5h5YtNXyD6WxUd4aSxIyWuGc", "https://www.youtube.com/watch?v=xNAm9O_duSA", "Who could be the next Prime Minister ?"]
|
81 |
+
],
|
82 |
+
)
|
83 |
+
|
84 |
+
iface.queue().launch()
|