CoolFace
Apppublic

bitpascals/bbb

sourceHugging Faceapache-2.0updated 1y agoView on Hugging Face
0likes
app.py131 linesDownload Raw Back to root
1# Import necessary libraries and modules2from youtube_transcript_api import YouTubeTranscriptApi3from youtube_transcript_api._errors import TranscriptsDisabled4from langchain.chains.combine_documents import create_stuff_documents_chain5from langchain.chains import create_retrieval_chain6from langchain_core.prompts import ChatPromptTemplate7from youtube_comment_downloader import YoutubeCommentDownloader8from langchain.text_splitter import RecursiveCharacterTextSplitter9from langchain_huggingface import HuggingFaceEmbeddings10from langchain_community.vectorstores import Chroma11from langchain_groq import ChatGroq12from dotenv import load_dotenv13import re14import os15import gradio as gr16 17# Load environment variables18load_dotenv()19groq_api_key = os.getenv("GROQ_API_KEY")20hf_token = os.getenv("HF_TOKEN")21 22if not groq_api_key:23    raise ValueError("Missing GROQ_API_KEY in environment variables")24 25os.environ["HF_TOKEN"] = hf_token or ""26 27# Set up models28llm = ChatGroq(api_key=groq_api_key, model_name="llama3-8b-8192", temperature=0.7)29 30embeddings = HuggingFaceEmbeddings(31    model_name="all-MiniLM-L6-v2",32    model_kwargs={'device': 'cpu'}33)34 35# Cache for storing processed video data36cache = {}37 38# Helper to extract video ID39def extract_video_id(url):40    pattern = r"(?:v=|youtu\.be/)([a-zA-Z0-9_-]{11})"41    match = re.search(pattern, url)42    return match.group(1) if match else None43 44# Gradio function45def query_video(video_url, query):46    try:47        if not video_url:48            return "Please provide a YouTube video URL."49 50        video_id = extract_video_id(video_url)51        if not video_id:52            return "Invalid YouTube URL format."53 54        # Use cached data if available55        if video_id in cache:56            vectors = cache[video_id]57        else:58            # Get transcript59            transcript = YouTubeTranscriptApi.get_transcript(video_id)60            transcript_texts = [entry['text'] for entry in transcript]61 62            # Get comments63            downloader = YoutubeCommentDownloader()64            comments = downloader.get_comments_from_url(video_url)65            comment_texts = [c['text'] for c in comments]66 67            if not transcript_texts and not comment_texts:68                return "No transcript or comments available."69 70            # Combine and split71            all_text = " ".join(transcript_texts + comment_texts)72            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=50)73            split_chunks = text_splitter.split_text(all_text)74 75            # Vector store76            vectors = Chroma.from_texts(77                texts=split_chunks,78                embedding=embeddings,79                persist_directory=f"chroma_db/{video_id}"80            )81            vectors.persist()82            cache[video_id] = vectors83 84        # Prompt85        prompt = ChatPromptTemplate.from_messages([86            ("system", 87            "You are a focused YouTube assistant. You have access to two sources: the video transcript (what was said) and the top viewer comments (opinions and reactions).\n\n"88            "Your job is to give short, clear, and accurate answers using ONLY the information provided. Do not guess or add anything not in the transcript or comments.\n\n"89            "Guidelines:\n"90            "1. Use the transcript for questions about what was said in the video.\n"91            "2. Use the comments for audience opinions.\n"92            "3. If both matter, combine them briefly.\n"93            "4. Mention timestamps (e.g., 'At 2:45...') for transcript quotes.\n"94            "5. Mention viewers (e.g., 'One comment said...') and likes if helpful.\n"95            "6. If the answer is not in the context, say: 'I couldn't find that in the transcript or comments.'\n"96            "7. Use short bullet points when listing things.\n"97            "8. If the word 'video' is in the question, use only the transcript and ignore comments."98            ),99            ("human", 100            "CONTEXT:\n{context}\n\n"101            "QUESTION:\n{input}\n\n"102            "Give a clear and concise answer using the rules above."103            )104        ])105 106        # Run retrieval107        document_chain = create_stuff_documents_chain(llm, prompt)108        retriever = vectors.as_retriever(search_kwargs={"k": 2})109        retrieval_chain = create_retrieval_chain(retriever, document_chain)110 111        if not query:112            return "Video loaded successfully! Ask me anything about it."113 114        result = retrieval_chain.invoke({"input": query})115        return result["answer"]116 117    except Exception as e:118        return f"Error: {str(e)}"119 120# Launch Gradio app121demo = gr.Interface(122    fn=query_video,123    inputs=["text", "text"],124    outputs="text",125    title="YouTube Transcript & Comment Analyzer",126    description="Paste a YouTube video URL and ask a question about the video or its comments."127)128 129if __name__ == "__main__":130    demo.launch()131