naiscriil/Youtube_comments

1import streamlit as st2import torch3import logging4import os5from model import load_resources, generate_comment6from youtube_api import get_prompt_for_model7 8logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')9logger = logging.getLogger(__name__)10 11if torch.cuda.is_available():12    DEVICE = "cuda"13    torch.cuda.empty_cache()14elif torch.backends.mps.is_available():15     DEVICE = "mps"16else:17    DEVICE = "cpu"18logger.info(f"Using device: {DEVICE}")19 20@st.cache_resource21def load_ml_resources():22    logger.info("Attempting to load model and tokenizer...")23    try:24        model, tokenizer = load_resources()25        if model is None or tokenizer is None:26             logger.error("load_resources returned None for model or tokenizer.")27             return None, None28        logger.info("Model and tokenizer loaded successfully via cache.")29        return model, tokenizer30    except Exception as e:31        logger.error(f"Fatal Error loading ML resources: {e}", exc_info=True)32        return None, None33 34model, tokenizer = load_ml_resources()35 36if model is None or tokenizer is None:37    st.error("Failed to load the language model or tokenizer. The application cannot start. Check logs for details.")38    st.stop()39 40st.set_page_config(page_title="YouTube Comment Generator", layout="wide")41st.title("🤖 YouTube Comment Generator")42st.markdown("Enter a YouTube video URL below and adjust the settings to generate a comment using a fine-tuned GPT-2 model.")43 44youtube_url = st.text_input("YouTube Video URL:", placeholder="e.g., https://www.youtube.com/watch?v=dQw4w9WgXcQ")45 46st.sidebar.header("⚙️ Generation Settings")47max_new_tokens = st.sidebar.slider("Max New Tokens:", min_value=10, max_value=200, value=75, step=5, help="Maximum number of tokens to generate for the comment.")48do_sample = st.sidebar.checkbox("Use Sampling:", value=True, help="Check to use probabilistic sampling (more creative), uncheck for deterministic greedy decoding (faster).")49 50if do_sample:51    temperature = st.sidebar.slider("Temperature:", min_value=0.1, max_value=2.0, value=0.9, step=0.05, help="Controls randomness. Lower values make output more focused, higher values more random.")52    top_k = st.sidebar.slider("Top-K Sampling:", min_value=0, max_value=100, value=50, step=1, help="Sample from the K most likely next tokens (0 disables).")53    top_p = st.sidebar.slider("Top-P (Nucleus) Sampling:", min_value=0.0, max_value=1.0, value=0.95, step=0.01, help="Sample from the smallest set of tokens whose cumulative probability exceeds P.")54else:55    temperature = 1.056    top_k = 057    top_p = 1.058 59no_repeat_ngram_size = st.sidebar.slider("No Repeat N-gram Size:", min_value=0, max_value=5, value=2, step=1, help="Prevents repeating sequences of N words (0 disables).")60repetition_penalty = st.sidebar.slider("Repetition Penalty:", min_value=1.0, max_value=2.0, value=1.0, step=0.05, help="Penalizes generating tokens that already appeared (1.0 = no penalty).")61 62 63if st.button("✨ Generate Comment"):64    if not youtube_url:65        st.warning("Please enter a YouTube URL.")66    else:67        with st.spinner("Fetching video details..."):68            logger.info(f"Fetching details for URL: {youtube_url}")69            prompt_text = get_prompt_for_model(youtube_url)70        if prompt_text:71            st.subheader("🧠 Generating Comment...")72            with st.spinner("Thinking..."):73                try:74                    generated_comment_text = generate_comment(75                        model=model,76                        tokenizer=tokenizer,77                        prompt=prompt_text,78                        max_new_tokens=max_new_tokens,79                        do_sample=do_sample,80                        top_k=top_k if do_sample else 0,81                        top_p=top_p if do_sample else 1.0,82                        temperature=temperature if do_sample else 1.0,83                        no_repeat_ngram_size=no_repeat_ngram_size,84                        repetition_penalty=repetition_penalty85                    )86                    logger.info(f"Raw generated text: {generated_comment_text}")87 88                    st.subheader("✅ Generated Comment:")89                    st.success(generated_comment_text)90 91                except Exception as e:92                    logger.error(f"Error during comment generation: {e}", exc_info=True)93                    st.error(f"An error occurred during generation: {e}")94 95        else:96            logger.warning(f"Could not retrieve details for URL: {youtube_url}")97            st.error(ERROR_MESSAGE if 'ERROR_MESSAGE' in globals() else "Error: Could not retrieve video details. Check URL and API key.")98 99else:100    st.info("Enter a YouTube URL and click 'Generate Comment'.")101 102st.sidebar.markdown("---")103st.sidebar.caption("Built with Streamlit and Hugging Face Transformers.")