File size: 4,609 Bytes
97c2b63 | 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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import os
import re
import gradio as gr
from youtube_transcript_api import YouTubeTranscriptApi, TranscriptsDisabled
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEndpointEmbeddings, HuggingFaceEndpoint, ChatHuggingFace
from langchain_community.vectorstores import FAISS
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
# No default token: user must supply their Hugging Face API token via the UI
def extract_video_id(url_or_id: str) -> str:
pattern = r"(?:v=|\/)([0-9A-Za-z_-]{11})"
match = re.search(pattern, url_or_id)
return match.group(1) if match else url_or_id
# Load, embed, and index the transcript
def load_vector_store(video_id: str, huggingface_token: str, embedding_model: str):
# Temporarily set the token for embedding calls
os.environ['HUGGINGFACEHUB_API_TOKEN'] = huggingface_token.strip()
try:
transcript_list = YouTubeTranscriptApi.get_transcript(video_id, languages=['en'])
transcript = ' '.join(chunk['text'] for chunk in transcript_list)
except TranscriptsDisabled:
transcript = ''
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
docs = splitter.create_documents([transcript])
embeddings = HuggingFaceEndpointEmbeddings(
model=embedding_model,
huggingfacehub_api_token=os.environ['HUGGINGFACEHUB_API_TOKEN']
)
return FAISS.from_documents(docs, embeddings)
# Initialize/reinitialize the QA chain
def setup(video_input, embedding_model, llm_model, huggingface_token):
video_id = extract_video_id(video_input)
vector_store = load_vector_store(video_id, huggingface_token, embedding_model)
retriever = vector_store.as_retriever(search_type='similarity', search_kwargs={'k': 4})
prompt_template = '''
You are a helpful assistant.
Answer ONLY from the provided transcript context.
If the context is insufficient, say you don't know.
{context}
Question: {question}
'''
prompt = PromptTemplate(template=prompt_template, input_variables=['context', 'question'])
memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)
# Configure the LLM endpoint
os.environ['HUGGINGFACEHUB_API_TOKEN'] = huggingface_token.strip()
hf_llm = HuggingFaceEndpoint(
repo_id=llm_model,
task='text-generation',
max_new_tokens=512,
temperature=0.2,
huggingfacehub_api_token=os.environ['HUGGINGFACEHUB_API_TOKEN']
)
chat_model = ChatHuggingFace(llm=hf_llm, verbose=True)
qa_chain = ConversationalRetrievalChain.from_llm(
llm=chat_model,
retriever=retriever,
memory=memory,
chain_type='stuff',
return_source_documents=False
)
# Reset chat history
return [], [], qa_chain
# Handle chat interactions
def respond(message, chat_history, qa_chain):
result = qa_chain({'question': message, 'chat_history': chat_history})
answer = result.get('answer') or result.get('result')
chat_history.append((message, answer))
return chat_history, chat_history
# Gradio UI layout
with gr.Blocks() as demo:
gr.Markdown('# Video Transcript Chatbot')
with gr.Row():
video_input = gr.Textbox(label='YouTube Video URL or ID', value='')
embedding_model_input = gr.Textbox(
label='Embedding Model (default: sentence-transformers/all-MiniLM-L6-v2)',
value='sentence-transformers/all-MiniLM-L6-v2'
)
llm_model_input = gr.Textbox(label='LLM Model Repo (e.g. google/flan-t5-large)', value='meta-llama/Llama-3.1-8B-Instruct')
token_input = gr.Textbox(label='Your HF API Token', placeholder='hf_...', type='password')
init_btn = gr.Button('Initialize Chat')
chatbot = gr.Chatbot()
chat_state = gr.State([])
chain_state = gr.State(None)
init_btn.click(
setup,
inputs=[video_input, embedding_model_input, llm_model_input, token_input],
outputs=[chatbot, chat_state, chain_state]
)
txt = gr.Textbox(placeholder='Ask a question about the video...', show_label=False)
txt.submit(respond, inputs=[txt, chat_state, chain_state], outputs=[chatbot, chat_state])
gr.Button('Clear Chat').click(lambda: ([], []), None, [chatbot, chat_state])
if __name__ == '__main__':
demo.launch() # pass share=True or host/port if needed
|