app.py

# Hey its your saiyan Utkarsh Shukla. I gonna write my custom comments after each line of code. So even a beginner (previous me) can read and get whats going on.
# Bear with me, I got your Back buddy
# And hey don't be overwhelmed by lines of code as most part is just covering design.


import streamlit as st  
# The saviour web app creator, easy peasy web app creation by few lines of codes. 
# No HTML, CSS, or JS needed!
# But I have used a bit to design but avoid it as its not necessary


from transformers import AutoTokenizer, AutoModelForCausalLM
# transformers here is just library which gives us access to Transformer architecture based pretrained models for natural language processing and other tasks. For now, Think of Transformer Architecture, as mystery Architecture which makes models way cooler.


# AutoTokenizer helps in Text input -> Sentences -> Words -> Even subwords like ['un', 'break', 'able'] -> Integer IDs that model expects.
# And whats awesome is Tokens will be generated following the configurations and requirements of model which we will be using.

# AutoModelForCausalLM is a powerful and convenient class serves as a high-level interface for loading pre-trained transformer models specifically designed for causal language modeling.
# The cooler part of these Auto* classes are you don't need to know exact class name of model(like GPT2LMHeadModel, CTRLLMHeadModel, ReformerLMHeadModel, etc.).
# Meaning The AutoModelForCausalLM automatically determine the correct model architecture based on the pretrained_model_name_or_path you provide. AutoModelForCausalLM infers this from the configuration files associated with the pre-trained model.


from wordcloud import WordCloud
# This will help us in knowing which words have large frequency. It creates a visual representation of words used, know as Word Cloud.
# More the frequency + More the importance -> Word will appear larger in Word Cloud.
# Mostly it avoids our stop words like it, is, are etc
# More frequency = more importance → bigger word in the cloud.


import matplotlib.pyplot as plt
# This guy helps us to plot. So wait till you see it.
# We’ll use it to show our Word Cloud in style.


import torch
# This import makes a library pyTorch available in our python code.
# This makes the PyTorch library available — a powerful math engine and deep learning framework our model runs on.
# Think of it as a toolkit which can do maths very very efficiently is being available for our code now.


@st.cache_resource
# A decorator in python is a way to enhance a function or a class. As they are followed by @ symbol
# The function above whome they are specified, the decorator code is executed both before and after of function code, on function call.
# Now here @st.cache_resource decorator is used before loading AutoTokenizer and AutoModelForCasualLM from gemm-2b.
# Cache the model and tokenizer to avoid reloading on every run
# So first run will load and save resources to global cache, and as user interact and causes rerun of load_model_and_tokenizer(), instead of loading again it will directly use cached resources from memory
def load_model_and_tokenizer():
    model_name = "google/gemma-2b"  # using gemma-2b for prototype for my GSOC Proposal. Wish me luck.
    #model_name = "openai-community/gpt2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    # Responsible for automatically downloading and loading the tokenizer configuration and vocabulary associated with the specified pre-trained model.
    # Downloads and loads the tokenizer config and vocab for the given model
    model = AutoModelForCausalLM.from_pretrained(model_name)
    # As we discussed, this class is designed for loading pre-trained language models for causal (next-token prediction) tasks.
    # Loads the actual model used for causal (next-word) prediction tasks
    return tokenizer, model


# Function to generate text with Gemma
# This calls the function we just made to get the tokenizer and model ready to work.
def generate_text(prompt, tone, max_length, temperature=0.7, top_p=0.9, repetition_penalty=1.0):
    tokenizer, model = load_model_and_tokenizer()
    # Adjust prompt based on tone
    # And Bro believe me, Impact: Better prompts = better outputs = a stronger GSoC impression!
    tone_prompts = {
        "Funny": f"Instruction: Generate a concise, humorous response to the following prompt. Prompt: {prompt}. Use witty wordplay, unexpected twists, or lighthearted exaggeration, avoiding offensive content. Aim for a punchline-style finish.",
        "Serious": f"Instruction: Provide a detailed, thoughtful, and professional response to the following prompt. Prompt: {prompt}. Offer logical reasoning, depth, and a formal tone, as if explaining to an expert audience.",
        "Poetic": f"Instruction: Write a vivid, poetic response to the following prompt. Prompt: {prompt}. Use metaphor, rhythm, and imagery to create a lyrical flow, as if crafting a short verse."
    }
    # This creates a dictionary that holds different prompts based on the tone we pick, making sure the model knows how to respond.
    input_text = tone_prompts.get(tone, prompt)
    # This picks the right instruction from the dictionary based on the tone.
    inputs = tokenizer(input_text, return_tensors="pt")
    input_ids = inputs["input_ids"]
    # This turns our input text (with the tone instruction) into a format (tensors) that the model can process using the tokenizer.
    input_token_length = input_ids.shape[1]  # Get the number of tokens in the input
    # Store the length of the input

    # --- Step 1: Estimate Tokens Needed (Increase the buffer) ---
    # Estimate slightly more tokens than words (e.g., 1.5x or 2x buffer)
    # Let's use a factor of 1.75 for a larger buffer to increase chances of reaching word count
    estimated_max_tokens = int(max_length * 1.75) 
    # Add a minimum token generation to avoid tiny requests
    estimated_max_tokens = max(estimated_max_tokens, 30) # Ensure we generate at least some tokens
    # --- Step 2: Generate with the higher token limit ---
    outputs = model.generate(
        inputs["input_ids"],
        # max_length=max_length + len(input_text.split()),
        # This sets how long the generated text can be. We add the number of words in our input text (len(input_text.split())) to the max_length the user picked, so the model knows how many total words to create.
        # CHANGE: Use max_new_tokens for clarity instead of calculating total length
        max_new_tokens = estimated_max_tokens, # Use the higher estimate, meaning 1.75 times the lenght
        # Generate THIS many NEW tokens
        temperature = temperature,
        # This controls how creative the model gets. A lower temperature (e.g., 0.7) keeps things more predictable, while a higher one makes it wilder and more random—think of it like adjusting the spice level!
        top_p=top_p,
        # This is like a filter for word choices. It picks from the top percentage of likely words (e.g., 0.9 means 90% of the best options), making the output diverse but not too crazy.
        repetition_penalty = repetition_penalty,
        # This stops the model from repeating the same words too much. A higher value (e.g., 1.5) pushes it to try new words, like telling it to mix up its vocabulary!
        num_return_sequences = 1,
        # This tells the model to give us just one version of the text. If we wanted more options, we could change
        do_sample = True,
        pad_token_id = tokenizer.eos_token_id # Good practice for generation
    )
    # --- Step 3: Decode ONLY the generated part ---
    # Slice the output tensor to get only the tokens AFTER the input tokens
    # This tells the model to generate text: it uses the input IDs, sets a max length, and adjusts creativity with temperature, top_p, and repetition_penalty.
    generated_token_ids = outputs[0, input_token_length:] 
    generated_text = tokenizer.decode(generated_token_ids, skip_special_tokens=True).strip()
    
    return generated_text # Return only the newly generated text
    # This turns again the model's output back into readable form, skipping any extra tokens we don’t need.
    
# Clean and Solid UI for our Project, keeping the blue theme of gemini.
# We will continue tutorial after this st.markdown()
st.markdown("""
    <style>
    /* Solid background with minimal decoration */
    .stApp {
        background-color: #1E1E3A;
        color: #f0f0f0;
    }
    
    /* Clean container styling for all elements */
    .stForm, div.stButton, .stMarkdown, div.stSlider {
        border-radius: 8px;
        transition: all 0.2s ease;
    }
    
    /* Professional title styling */
    .title {
        color: #ffffff;
        font-size: 32px;
        font-weight: 700;
        padding: 15px 0;
        font-family: 'Segoe UI', Arial, sans-serif;
        letter-spacing: 0.5px;
        margin-bottom: 15px;
        border-bottom: 2px solid #4A5BEA;
        display: inline-block;
    }
    
    /* Clean instructions card */
    .instructions {
        background: rgba(255, 255, 255, 0.05);
        padding: 20px;
        border-radius: 8px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        font-size: 16px;
        color: #f0f0f0;
        border: 1px solid rgba(255, 255, 255, 0.1);
        margin-bottom: 20px;
        line-height: 1.5;
    }
    
    /* Professional output box */
    .output-box {
        background: rgba(20, 20, 40, 0.8);
        padding: 20px;
        border-radius: 8px;
        box-shadow: 0 4px 8px rgba(0,0,0,0.2);
        font-family: 'Courier New', monospace;
        font-size: 16px;
        color: #ffffff;
        white-space: pre-wrap;
        border: 1px solid #4A5BEA;
        margin-top: 1.5rem;
    }
    
    /* Form styling */
    .stTextInput>div>div>input {
        background: rgba(255, 255, 255, 0.05);
        border: 1px solid rgba(255, 255, 255, 0.2);
        color: white;
        border-radius: 6px;
        padding: 10px;
        font-size: 16px;
    }
    .stTextInput>div>div>input:focus {
        border-color: #4A5BEA;
        box-shadow: 0 0 5px rgba(74, 91, 234, 0.5);
    }
    
    /* Clean, professional buttons */
    .stButton>button {
        background-color: #4A5BEA;
        color: white;
        border: none;
        border-radius: 6px;
        padding: 10px 15px;
        font-size: 16px;
        font-weight: 500;
        box-shadow: 0 2px 5px rgba(0,0,0,0.2);
        transition: all 0.2s;
        text-transform: none;
        letter-spacing: 0.5px;
    }
    .stButton>button:hover {
        background-color: #3A4AC0;
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.2);
    }
    
    /* Example button styling */
    .example-button {
        background-color: #2D3250;
        color: white;
        border: 1px solid #4A5BEA;
        border-radius: 6px;
        padding: 8px 16px;
        margin: 5px;
        cursor: pointer;
        transition: all 0.2s;
        font-size: 14px;
        font-weight: 500;
    }
    .example-button:hover {
        background-color: #4A5BEA;
    }
    
    /* Enhanced sliders */
    .stSlider>div>div>div {
        background: #4A5BEA !important;
    }
    
    /* Improved slider number values styling */
    .stSlider p {
        color: #f0f0f0 !important;
    }
    
    /* Fix for the blue box numbers on sliders */
    .stSlider > div > div > div > div > div > div {
        background-color: #4A5BEA !important;
        color: white !important;
        font-weight: 500 !important;
        border: none !important;
        border-radius: 4px !important;
        padding: 2px 6px !important;
        font-size: 14px !important;
    }
    
    /* Number value containers at the ends of sliders */
    .stSlider > div > div > div:first-child,
    .stSlider > div > div > div:last-child {
        background-color: #292952 !important;
        color: #f0f0f0 !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        border-radius: 4px !important;
        padding: 4px 8px !important;
        font-size: 14px !important;
        font-family: monospace !important;
    }
    
    /* Custom selectbox styling */
    .stSelectbox>div>div {
        background: rgba(255, 255, 255, 0.05);
        border: 1px solid rgba(255, 255, 255, 0.2);
        border-radius: 6px;
    }
    
    /* Expander styling */
    .streamlit-expanderHeader {
        background: rgba(255, 255, 255, 0.05);
        border-radius: 6px;
        border: 1px solid rgba(255, 255, 255, 0.1);
        color: #f0f0f0 !important;
        font-weight: 500;
    }
    
    /* WordCloud container */
    .wordcloud-container {
        background: rgba(255, 255, 255, 0.05);
        border-radius: 8px;
        padding: 15px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        margin-top: 20px;
        border: 1px solid rgba(255, 255, 255, 0.1);
    }
    
    /* Footer styling */
    .footer {
        background: rgba(0, 0, 0, 0.2);
        border-radius: 8px;
        padding: 15px;
        margin-top: 30px;
        border-top: 1px solid rgba(255, 255, 255, 0.1);
        text-align: center;
    }
    
    /* Parameter card styling */
    .parameter-card {
        background: rgba(255, 255, 255, 0.05);
        border-radius: 8px;
        padding: 15px;
        margin-bottom: 15px;
        border: 1px solid rgba(255, 255, 255, 0.1);
    }
    .parameter-card h4 {
        color: #4A5BEA;
        margin-top: 0;
        font-size: 16px;
        font-weight: 500;
    }
    
    /* Scrollbar styling */
    ::-webkit-scrollbar {
        width: 8px;
        height: 8px;
    }
    ::-webkit-scrollbar-track {
        background: rgba(255, 255, 255, 0.05);
        border-radius: 4px;
    }
    ::-webkit-scrollbar-thumb {
        background: #4A5BEA;
        border-radius: 4px;
    }
    ::-webkit-scrollbar-thumb:hover {
        background: #3A4AC0;
    }
    </style>
""", unsafe_allow_html=True)

# Its our Header with GSoC logo, keeping simple as deadline is coming. :)
st.markdown("""
    <div style="display: flex; align-items: center; justify-content: space-between; margin-bottom: 25px;">
        <h1 class="title">Gemma Text Generator</h1>
        <img src="https://huggingface.co/spaces/Kakaarot/Gemma-HuggingFace_TextCompletion_Demo/resolve/main/images/GSoC_logo.png" width="100" alt="GSoC 2025" style="margin-top: 5px;">
    </div>
""", unsafe_allow_html=True)


# This part is generally telling how things work
st.markdown("""
    <div class="instructions">
        <b><a href = "https://huggingface.co/spaces/Kakaarot/Gemma-HuggingFace_TextCompletion_Demo/discussions/1">Please check the discussion</a><b>, I mentioned there the reason, why your first response will take little more time. 
        Thanks for understanding, Now Enjoyyy 😁 <br><br>
        Enter a prompt below to generate text using the Gemma model from DeepMind. Customize the tone and length to see different outputs!<br><br>
        <b>Example:</b> Prompt: "The cat sat on" | Tone: "Funny" | Length: 50 → "The cat sat on my homework and laughed as I cried over my grades."
    </div>
""", unsafe_allow_html=True)


# While this tells specifically step by step
with st.expander("How does this work?"):
    st.markdown("""
    <div style="padding: 10px;">
        <ul>
            <li style="margin-bottom: 8px;">This app uses <b>Gemma-2B</b>, a language model from Google DeepMind.</li>
            <li style="margin-bottom: 8px;">You give it a prompt, and it predicts the next words one-by-one (aka causal language modeling).</li>
            <li style="margin-bottom: 8px;">The <b>tone</b> you choose adds flavor to the prompt before it hits the model.</li>
            <li style="margin-bottom: 8px;">Parameters like <b>temperature</b> control how wild or safe the answers are.</li>
            <li>The output is visualized in a <b>Word Cloud</b> so you can see which words stand out!</li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

# Clean example buttons area
st.markdown("<p style='margin-top: 20px; margin-bottom: 10px;'>Try these examples:</p>", unsafe_allow_html=True)
# Below we defined one click example, where one click will fill the form generate button will be pressed automatically,
# Fixed example buttons layout
col1, col2, col3 = st.columns(3)

# Hey! This is for giving user a taste of what would happen, so that he can try out the process on one click
if "trigger_example" not in st.session_state:
    st.session_state.trigger_example = False

with col1:
    if st.button("✨ Funny Cat Story"):
        st.session_state.prompt = "The cat hacked my WiFi"
        # This sets the prompt to a fun example when the button is clicked
        st.session_state.tone = "Funny"
        # This sets the tone to Funny for the example.
        st.session_state.trigger_example = True
        # This turns on the trigger to automatically generate the example.
        # And then same things with below Poetic and Serious tones of ours, Phewww, don't be sleepy, little more.
with col2:
    if st.button("🌅 Poetic Goodbye"):
        st.session_state.prompt = "As the sun set on our final day"
        st.session_state.tone = "Poetic"
        st.session_state.trigger_example = True
with col3:
    if st.button("🧠 Serious AI Future"):
        st.session_state.prompt = "The future of AI is"
        st.session_state.tone = "Serious"
        st.session_state.trigger_example = True

# Clean form with better spacing
with st.form(key="input_form"):
    # This starts a form where users can input their own prompt and settings
    st.markdown('<div style="margin-bottom: 15px;"><h3 style="color: #4A5BEA; margin-bottom: 10px;">Generate Your Text</h3></div>', unsafe_allow_html=True)
    # This adds a styled heading to label the form section.
    prompt = st.text_input("Enter a prompt", placeholder="e.g., 'The future of AI is'", value=st.session_state.get("prompt", ""))
    # This creates a text box where users can type their prompt, with a placeholder hint and a default value from our example if set.
    col1, col2 = st.columns(2)
    # This splits the next part into two columns for better layout. Yep Saiyans like orgainised layouts.
    with col1:
        tone = st.selectbox("Tone", ["Funny", "Serious", "Poetic"], index=["Funny", "Serious", "Poetic"].index(st.session_state.get("tone", "Funny")))
    with col2:
        max_length = st.slider("Word count", 20, 100, 50, help="Tries to generate text close to this word count. Output might be shorter if the model finishes early, or slightly different due to word splitting. I am considering 1.75 tokens as one word.")
        # This adds a slider for users to set how many words they want in the output, ranging from 20 to 100 with a default of 50.
        # And similarly every slider here works
    
    st.markdown('<div class="parameter-card"><h4>Advanced Parameters</h4>', unsafe_allow_html=True)
    
    # Adding custom number formatting to make slider values look better
    st.markdown("""
    <style>
    /* Custom number styling for slider values */
    .value-display {
        background-color: #292952;
        padding: 4px 8px;
        border-radius: 4px;
        font-family: monospace;
        border: 1px solid rgba(255, 255, 255, 0.1);
        font-size: 14px;
        display: inline-block;
        margin-left: 5px;
        color: #f0f0f0;
    }
    </style>
    """, unsafe_allow_html=True)
    
    col1, col2 = st.columns(2)
    with col1:
        temperature = st.slider("Temperature (Creativity)", 0.2, 1.5, 0.7, help="Higher values make output more random")
    with col2:
        top_p = st.slider("Top-p (Nucleus Sampling)", 0.1, 1.0, 0.9, help="Controls diversity")
    
    repetition_penalty = st.slider("Repetition Penalty", 1.0, 2.0, 1.0, help="Higher values discourage repetition")
    
    st.markdown('</div>', unsafe_allow_html=True)
    
    submit_button = st.form_submit_button(label="Generate")

# Generate and display output
if submit_button or st.session_state.trigger_example:
# This checks if the Generate button was clicked or our predefined one click example was triggered.
    st.session_state.trigger_example = False
    # This resets the example trigger so it doesn’t keep running.
    if not prompt:
        st.error("Please enter a prompt!")
        # If user thought clicking generate is fun without entering a prompt. Naah buddy, i stopped you :)
    else:
        with st.spinner("Generating text..."):
            # This shows a spinning icon while the text is being created.
            output = generate_text(prompt, tone, max_length, temperature, top_p, repetition_penalty)
            
            # Display metadata about the generation with improved value styling
            st.markdown(f"""
            <div style="background: rgba(255, 255, 255, 0.05); border-radius: 8px; padding: 12px; 
                margin-bottom: 15px; font-size: 14px; border: 1px solid rgba(255, 255, 255, 0.1);">
                <span style="color: #4A5BEA; font-weight: 500;">Tone:</span> 
                <span class="value-display">{tone}</span> | 
                <span style="color: #4A5BEA; font-weight: 500;">Temperature:</span> 
                <span class="value-display">{temperature:.2f}</span> | 
                <span style="color: #4A5BEA; font-weight: 500;">Words:</span> 
                <span class="value-display">~{max_length}</span>
            </div>
            """, unsafe_allow_html=True)
            
            st.markdown(f'<div class="output-box">{output}</div>', unsafe_allow_html=True)
            
            # WordCloud visualization
            st.markdown('<div class="wordcloud-container">', unsafe_allow_html=True)
            st.markdown('<h4 style="color: #4A5BEA; margin-top: 0;">Word Cloud Visualization</h4>', unsafe_allow_html=True)
            # This starts a container for the word cloud visualization.
            
            # Generate a clean wordcloud
            wordcloud = WordCloud(
                width=600, 
                height=300, 
                background_color="#1E1E3A",
                colormap="viridis",
                max_words=100,
                contour_width=0
            ).generate(output)
            
            plt.figure(figsize=(10, 5))
            plt.imshow(wordcloud, interpolation="bilinear")
            plt.axis("off")
            st.pyplot(plt)
            st.markdown('</div>', unsafe_allow_html=True)

# Clean footer
st.markdown("""
<div class="footer">
    <p>Built with ❤️ by Utkarsh Shukla for GSoC Proposal 2025 | Powered by Gemma + Hugging Face</p>
    <p style="margin-top: 10px;">
        <img src="https://huggingface.co/spaces/Kakaarot/Gemma-HuggingFace_TextCompletion_Demo/resolve/main/images/google-gemini-icon.png" width="100" alt="Gemma by DeepMind">
    </p>
    <p style="margin-top: 10px;">Wish me luck! 🤞</p>         
</div>
""", unsafe_allow_html=True)

# Yep here we done. Hope you guys like my attempt. I really enjoyed this project as Google Gen AI + Kaggle Workshop also helped. Thanks for your time
# And please let me know the improvements i need, Will appreciate any reviews.