app.py

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import spaces # Required for Hugging Face Spaces GPU
import random
import numpy as np

# Set random seeds for reproducibility
def set_random_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# Set initial seed
set_random_seed(42)

# Determine the device to use (GPU if available, otherwise CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "FlameF0X/SnowflakeCore-G1-Tiny2",
    trust_remote_code=True,
    force_download=True,
    use_safetensors=True,
).to(device)

tokenizer = AutoTokenizer.from_pretrained(
    "FlameF0X/SnowflakeCore-G1-Tiny2",
    trust_remote_code=True,
    force_download=True,
    use_safetensors=True,
)

@spaces.GPU # Required decorator for GPU usage in Hugging Face Spaces
def advanced_generate(prompt, max_length=50, temperature=1.0, top_k=50, top_p=0.9, 
                      repetition_penalty=1.1, do_sample=True, seed=None):
    """
    Generates text with advanced sampling parameters.
    The model and input tensors are moved to the appropriate device (GPU/CPU).
    """
    # Set seed if provided
    if seed is not None:
        set_random_seed(seed)
    
    model.eval()
    # Move input_ids to the same device as the model
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    generated = input_ids
    
    with torch.no_grad():
        for i in range(max_length):
            # Get model outputs
            outputs = model(input_ids=generated)
            next_token_logits = outputs["logits"][:, -1, :]
            
            # Apply repetition penalty
            if repetition_penalty != 1.0:
                for token_id in set(generated[0].tolist()):
                    next_token_logits[0, token_id] /= repetition_penalty
            
            # Apply temperature
            if temperature != 1.0:
                next_token_logits = next_token_logits / temperature
            
            # Convert logits to probabilities
            probs = torch.softmax(next_token_logits, dim=-1)
            
            if do_sample and temperature > 0:
                # Apply top-k filtering
                if top_k > 0:
                    top_k_probs, top_k_indices = torch.topk(probs, min(top_k, probs.size(-1)))
                    probs_filtered = torch.zeros_like(probs)
                    probs_filtered.scatter_(1, top_k_indices, top_k_probs)
                    probs = probs_filtered
                
                # Apply top-p (nucleus) filtering
                if top_p < 1.0:
                    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
                    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
                    
                    # Remove tokens with cumulative probability above the threshold
                    sorted_indices_to_remove = cumulative_probs > top_p
                    # Keep at least one token
                    sorted_indices_to_remove[0, 0] = False
                    
                    # Create mask for tokens to remove
                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                    probs[indices_to_remove] = 0
                    
                    # Renormalize
                    probs = probs / probs.sum(dim=-1, keepdim=True)
                
                # Sample from the filtered distribution
                next_token_id = torch.multinomial(probs, num_samples=1)
            else:
                # Greedy decoding
                next_token_id = torch.argmax(probs, dim=-1).unsqueeze(-1)
            
            generated = torch.cat((generated, next_token_id), dim=1)
            
            # Check for end of sequence
            if next_token_id.item() == tokenizer.eos_token_id:
                break
                
    return tokenizer.decode(generated[0], skip_special_tokens=True)

def gradio_generate(prompt, max_length, temperature, top_k, top_p, repetition_penalty, do_sample, seed):
    """
    Wrapper function for Gradio interface.
    """
    # Convert seed to int if provided, otherwise use random
    if seed == "" or seed is None:
        seed = random.randint(0, 2**32 - 1)
    else:
        try:
            seed = int(seed)
        except ValueError:
            seed = random.randint(0, 2**32 - 1)
    
    # Ensure parameters are in valid ranges
    max_length = max(1, min(200, int(max_length)))
    temperature = max(0.1, min(2.0, float(temperature)))
    top_k = max(1, min(100, int(top_k)))
    top_p = max(0.1, min(1.0, float(top_p)))
    repetition_penalty = max(1.0, min(2.0, float(repetition_penalty)))
    
    return advanced_generate(
        prompt=prompt,
        max_length=max_length,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=do_sample,
        seed=seed
    )

# Custom CSS for ultra-compact UI and full width
custom_css = """
.gradio-container {
    max-width: 100% !important; /* Changed from 1200px to 100% for full width */
    padding: 10px !important;
}
.compact-slider {
    margin: 2px 0 !important;
}
.parameter-info {
    font-size: 10px;
    color: #888;
    margin: -8px 0 8px 0 !important;
    line-height: 1.1;
    padding: 2px 4px;
    background: rgba(0,0,0,0.05);
    border-radius: 3px;
}
.gradio-group {
    padding: 8px !important;
    margin: 4px 0 !important;
}
.gradio-textbox {
    margin-bottom: 8px !important;
}
.gradio-slider {
    margin: 4px 0 !important;
}
.gradio-checkbox {
    margin: 4px 0 !important;
}
.gradio-number {
    margin: 4px 0 !important;
}
.compact-header {
    margin: 0 0 10px 0 !important;
    padding: 8px !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 8px;
    color: white;
}
.control-section {
    background: rgba(255,255,255,0.02);
    border-radius: 6px;
    padding: 8px !important;
    margin: 4px 0 !important;
}
"""

# Create the Gradio interface with ultra-compact layout
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="SnowflakeCore Text Generator") as iface:
    
    # Compact header
    gr.HTML(f"""
    <div class="compact-header">
        <h2 style="margin: 0; font-size: 18px;">🔥 SnowflakeCore-G1-Tiny2 | Running on: {device}</h2>
    </div>
    """)
    
    with gr.Row():
        # Left column - Input/Output (60% width)
        with gr.Column(scale=6, min_width=400):
            prompt_input = gr.Textbox(
                lines=4, 
                placeholder="Enter your prompt here...", 
                label="📝 Input Prompt",
                show_label=True
            )
            
            generate_btn = gr.Button("🚀 Generate", variant="primary", size="sm", scale=1)
            
            output_text = gr.Textbox(
                lines=8, 
                label="✨ Generated Text",
                show_label=True,
                interactive=False
            )
        
        # Right column - Parameters (40% width)  
        with gr.Column(scale=4, min_width=300):
            gr.HTML("<div style='font-weight: bold; font-size: 14px; margin-bottom: 8px; color: #333;'>⚙️ Parameters</div>")
            
            # Core parameters in compact group
            with gr.Group(elem_classes=["control-section"]):
                max_length = gr.Slider(10, 2048, value=100, step=5, label="Max Length", elem_classes=["compact-slider"])
                gr.HTML("<div class='parameter-info'>📏 Tokens to generate (10-2048)</div>")
                
                temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature", elem_classes=["compact-slider"])
                gr.HTML("<div class='parameter-info'>🌡️ Creativity: 0.1=focused, 2.0=creative</div>")
            
            # Advanced parameters  
            with gr.Group(elem_classes=["control-section"]):
                top_k = gr.Slider(1, 150, value=50, step=1, label="Top-K", elem_classes=["compact-slider"])
                gr.HTML("<div class='parameter-info'>🎯 Word choice diversity</div>")
                
                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P", elem_classes=["compact-slider"])
                gr.HTML("<div class='parameter-info'>🎪 Probability cutoff</div>")
                
                repetition_penalty = gr.Slider(1.0, 3.0, value=1.15, step=0.05, label="Rep. Penalty", elem_classes=["compact-slider"])
                gr.HTML("<div class='parameter-info'>🔄 Anti-repetition strength</div>")
            
            # Controls row
            with gr.Row():
                do_sample = gr.Checkbox(value=True, label="Sampling", scale=1)
                seed_input = gr.Number(label="Seed", value=None, precision=0, scale=1, minimum=0)
            gr.HTML("<div class='parameter-info'>🎲 Sampling=creative | Seed=reproducible</div>")
    
    # Examples section
    gr.HTML("<h3 style='margin-top: 25px; margin-bottom: 10px;'>💡 Quick Examples</h3>")
    
    examples = gr.Examples(
        examples=[
            ["Once upon a time in a magical forest,", 120, 0.8, 40, 0.9, 1.2, True, 42],
            ["The future of artificial intelligence is", 80, 1.0, 50, 0.95, 1.1, True, None],
            ["In a world where technology and nature coexist,", 150, 1.2, 60, 0.85, 1.3, True, 123],
            ["Write a haiku about winter:", 50, 0.7, 30, 0.8, 1.0, True, None],
            ["Explain quantum computing in simple terms:", 200, 0.6, 40, 0.9, 1.1, True, None]
        ],
        inputs=[prompt_input, max_length, temperature, top_k, top_p, repetition_penalty, do_sample, seed_input],
        outputs=[output_text],
        fn=gradio_generate,
        cache_examples=False,
        label=None
    )
    
    # Event handlers
    generate_btn.click(
        fn=gradio_generate,
        inputs=[prompt_input, max_length, temperature, top_k, top_p, repetition_penalty, do_sample, seed_input],
        outputs=[output_text]
    )
    
    # Tips section
    with gr.Accordion("📚 Parameter Guide & Tips", open=False):
        gr.HTML("""
        <div style="font-size: 12px; line-height: 1.4;">
        <h4>🎛️ Parameter Combinations for Different Use Cases:</h4>
        
        <strong>📝 Creative Writing:</strong> Temperature: 0.8-1.2, Top-K: 40-60, Top-P: 0.85-0.95, Rep. Penalty: 1.1-1.3<br>
        <strong>📋 Factual/Technical:</strong> Temperature: 0.3-0.7, Top-K: 20-40, Top-P: 0.9-1.0, Rep. Penalty: 1.0-1.1<br>
        <strong>🎭 Experimental/Artistic:</strong> Temperature: 1.2-2.0, Top-K: 60-100, Top-P: 0.7-0.9, Rep. Penalty: 1.2-1.5<br>
        <strong>🎯 Focused/Consistent:</strong> Temperature: 0.1-0.5, Top-K: 10-30, Top-P: 0.8-0.95, Rep. Penalty: 1.0-1.2<br><br>
        
        <strong>💡 Pro Tips:</strong><br>
        • Use same seed for reproducible results across generations<br>
        • Higher repetition penalty helps with stuck loops or repeated phrases<br>
        • Lower temperature + higher top-p = focused but varied vocabulary<br>
        • Disable sampling for completely deterministic output (useful for testing)<br>
        • Start with defaults and adjust one parameter at a time to understand effects
        </div>
        """)

# Launch the Gradio application
if __name__ == "__main__":
    iface.launch()