SandyBot commited on
Commit
9f3be5a
·
verified ·
1 Parent(s): c50146d

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +37 -0
  2. rag_chain.py +92 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from rag_chain import get_rag_chain
3
+
4
+ st.set_page_config(page_title="Cricket Chat", page_icon="🏏")
5
+ st.title("YouTube Video Chatbot")
6
+
7
+ video_id = st.text_input("Enter YouTube video ID (e.g. Gfr50f6ZBvo)")
8
+
9
+ if video_id and st.session_state.get("current_video") != video_id:
10
+ with st.spinner("Processing transcript..."):
11
+ try:
12
+ st.session_state.chain = get_rag_chain(video_id)
13
+ st.session_state.current_video = video_id
14
+ st.session_state.messages = []
15
+ st.success("Ready! Ask away.")
16
+ except ValueError as e:
17
+ st.error(str(e))
18
+ st.stop()
19
+
20
+ if "messages" not in st.session_state:
21
+ st.session_state.messages = []
22
+
23
+ for msg in st.session_state.messages:
24
+ with st.chat_message(msg["role"]):
25
+ st.markdown(msg["content"])
26
+
27
+ if "chain" in st.session_state:
28
+ if prompt := st.chat_input("Ask about the video..."):
29
+ st.session_state.messages.append({"role": "user", "content": prompt})
30
+ with st.chat_message("user"):
31
+ st.markdown(prompt)
32
+
33
+ with st.chat_message("assistant"):
34
+ with st.spinner("Thinking..."):
35
+ response = st.session_state.chain.invoke(prompt)
36
+ st.markdown(response)
37
+ st.session_state.messages.append({"role": "assistant", "content": response})
rag_chain.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from youtube_transcript_api import YouTubeTranscriptApi, TranscriptsDisabled
3
+ from langchain_text_splitters import RecursiveCharacterTextSplitter
4
+ from langchain_community.vectorstores import FAISS
5
+ from langchain_core.prompts import PromptTemplate
6
+ from langchain_huggingface import HuggingFaceEmbeddings
7
+ from langchain_groq import ChatGroq
8
+ from langchain_core.output_parsers import StrOutputParser
9
+ from langchain_core.runnables import RunnableParallel, RunnablePassthrough, RunnableLambda
10
+
11
+ VECTORSTORE_DIR = "vectorstores"
12
+ EMBEDDING_MODEL_NAME = os.environ["EMBEDDING_MODEL"]
13
+
14
+ _embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
15
+
16
+
17
+ def load_transcript(video_id: str) -> str:
18
+ """Fetch and flatten a YouTube video's transcript."""
19
+ try:
20
+ yt_api = YouTubeTranscriptApi()
21
+ transcript_list = yt_api.fetch(video_id, languages=["en"])
22
+ transcript = " ".join(chunk.text for chunk in transcript_list)
23
+ return transcript
24
+ except TranscriptsDisabled:
25
+ raise ValueError("No captions available for this video.")
26
+
27
+
28
+ def split_transcript(transcript: str):
29
+ splitter = RecursiveCharacterTextSplitter(
30
+ chunk_size=1000,
31
+ chunk_overlap=200,
32
+ )
33
+ return splitter.create_documents([transcript])
34
+
35
+
36
+ def get_or_build_vectorstore(video_id: str):
37
+ """Load a cached FAISS index for this video, or build + save one if missing."""
38
+ path = os.path.join(VECTORSTORE_DIR, video_id)
39
+
40
+ if os.path.exists(path):
41
+ return FAISS.load_local(
42
+ path,
43
+ _embedding_model,
44
+ allow_dangerous_deserialization=True,
45
+ )
46
+
47
+ transcript = load_transcript(video_id)
48
+ chunks = split_transcript(transcript)
49
+ vectorstore = FAISS.from_documents(chunks, _embedding_model)
50
+
51
+ os.makedirs(VECTORSTORE_DIR, exist_ok=True)
52
+ vectorstore.save_local(path)
53
+
54
+ return vectorstore
55
+
56
+
57
+ def format_docs(retrieved_docs):
58
+ return "\n\n".join(doc.page_content for doc in retrieved_docs)
59
+
60
+
61
+ def get_rag_chain(video_id: str):
62
+ """Builds and returns the full RAG chain for a given YouTube video."""
63
+ vectorstore = get_or_build_vectorstore(video_id)
64
+ retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
65
+
66
+ llm = ChatGroq(
67
+ model="llama-3.1-8b-instant",
68
+ temperature=0,
69
+ api_key=os.environ["GROQ_API_KEY"],
70
+ )
71
+
72
+ prompt = PromptTemplate(
73
+ template="""
74
+ You are a helpful assistant.
75
+ Answer ONLY from the provided transcript context.
76
+ If the context is insufficient, just say you don't know.
77
+
78
+ {context}
79
+ Question: {question}
80
+ """,
81
+ input_variables=["context", "question"],
82
+ )
83
+
84
+ parser = StrOutputParser()
85
+
86
+ parallel_chain = RunnableParallel({
87
+ "context": retriever | RunnableLambda(format_docs),
88
+ "question": RunnablePassthrough(),
89
+ })
90
+
91
+ main_chain = parallel_chain | prompt | llm | parser
92
+ return main_chain