app.py

import os
import torch
import gradio as gr
import datetime
from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
from safetensors.torch import load_file

# Constants
MODEL_CONFIG = {
    "G0-Release": "FlameF0X/Snowflake-G0-Release",
    "G0-Release-2": "FlameF0X/Snowflake-G0-Release-2",
    "G0-Release-2.5": "FlameF0X/Snowflake-G0-Release-2.5"
}

MAX_LENGTH = 384
TEMPERATURE_DEFAULT = 0.7
TOP_P_DEFAULT = 0.9
TOP_K_DEFAULT = 40
MAX_NEW_TOKENS_DEFAULT = 256

# UI parameter bounds
TEMPERATURE_MIN, TEMPERATURE_MAX = 0.1, 2.0
TOP_P_MIN, TOP_P_MAX = 0.1, 1.0
TOP_K_MIN, TOP_K_MAX = 1, 100
MAX_NEW_TOKENS_MIN, MAX_NEW_TOKENS_MAX = 16, 1024

# Styling
css = """
.gradio-container { background-color: #1e1e2f !important; color: #e0e0e0 !important; }
.header { background-color: #2b2b3c; padding: 20px; margin-bottom: 20px; border-radius: 10px; text-align: center; }
.header h1 { color: #66ccff; margin-bottom: 10px; }
.snowflake-icon { font-size: 24px; margin-right: 10px; }
.footer { text-align: center; margin-top: 20px; font-size: 0.9em; color: #999; }
.parameter-section { background-color: #2a2a3a; padding: 15px; border-radius: 8px; margin-bottom: 15px; }
.parameter-section h3 { margin-top: 0; color: #66ccff; }
.example-section { background-color: #223344; padding: 15px; border-radius: 8px; margin-bottom: 15px; }
.example-section h3 { margin-top: 0; color: #66ffaa; }
.model-select { background-color: #2a2a4a; padding: 10px; border-radius: 8px; margin-bottom: 15px; }
"""

# Model registry
model_registry = {}

def load_all_models():
    for name, model_id in MODEL_CONFIG.items():
        print(f"Loading model: {name} from {model_id}")
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        safetensor_path = os.path.join(model_id, "model.safetensors")
        if os.path.exists(safetensor_path):
            print("Loading from safetensors...")
            model = load_file(safetensor_path)
        else:
            print("Loading from Hugging Face or .bin...")
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto"
            )

        pipeline = TextGenerationPipeline(
            model=model,
            tokenizer=tokenizer,
            return_full_text=False,
            max_length=MAX_LENGTH
        )

        model_registry[name] = (model, tokenizer, pipeline)

def generate_text(prompt, model_version, temperature, top_p, top_k, max_new_tokens, history=None):
    if history is None:
        history = []
    history.append({"role": "user", "content": prompt})
    
    try:
        if model_version not in model_registry:
            raise ValueError(f"Model '{model_version}' not found.")

        _, tokenizer, pipeline = model_registry[model_version]

        outputs = pipeline(
            prompt,
            do_sample=temperature > 0,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            max_new_tokens=max_new_tokens,
            pad_token_id=tokenizer.pad_token_id,
            num_return_sequences=1
        )

        response = outputs[0]["generated_text"]
        history.append({"role": "assistant", "content": response, "model": model_version})

        formatted_history = []
        for entry in history:
            prefix = "👤 User: " if entry["role"] == "user" else f"❄️ [{entry.get('model', 'Model')}]: "
            formatted_history.append(f"{prefix}{entry['content']}")

        return response, history, "\n\n".join(formatted_history)
    
    except Exception as e:
        error_msg = f"Error generating response: {str(e)}"
        history.append({"role": "assistant", "content": f"[ERROR] {error_msg}", "model": model_version})
        return error_msg, history, str(history)

def clear_conversation():
    return "", [], ""

def create_demo():
    with gr.Blocks(css=css) as demo:
        # Header
        gr.HTML("""
        <div class="header">
            <h1><span class="snowflake-icon">❄️</span> Snowflake Models Demo</h1>
            <p>Experience the capabilities of the Snowflake series language models</p>
        </div>
        """)
        
        with gr.Column():
            with gr.Row(elem_classes="model-select"):
                model_version = gr.Radio(
                    choices=list(MODEL_CONFIG.keys()),
                    value=list(MODEL_CONFIG.keys())[0],
                    label="Select Model Version",
                    info="Choose which Snowflake model to use"
                )
            
            chat_history_display = gr.Textbox(
                value="", 
                label="Conversation History", 
                lines=10, 
                max_lines=30,
                interactive=False
            )
            
            history_state = gr.State([])
            
            with gr.Row():
                with gr.Column(scale=4):
                    prompt = gr.Textbox(
                        placeholder="Type your message here...", 
                        label="Your Input",
                        lines=2
                    )
                with gr.Column(scale=1):
                    submit_btn = gr.Button("Send", variant="primary")
                    clear_btn = gr.Button("Clear Conversation")

            response_output = gr.Textbox(
                value="", 
                label="Model Response", 
                lines=5,
                max_lines=10,
                interactive=False
            )
        
        # Generation Parameters
        with gr.Accordion("Generation Parameters", open=False):
            with gr.Column(elem_classes="parameter-section"):
                with gr.Row():
                    with gr.Column():
                        temperature = gr.Slider(
                            minimum=TEMPERATURE_MIN, maximum=TEMPERATURE_MAX,
                            value=TEMPERATURE_DEFAULT, step=0.05,
                            label="Temperature"
                        )
                        top_p = gr.Slider(
                            minimum=TOP_P_MIN, maximum=TOP_P_MAX,
                            value=TOP_P_DEFAULT, step=0.05,
                            label="Top-p"
                        )
                    with gr.Column():
                        top_k = gr.Slider(
                            minimum=TOP_K_MIN, maximum=TOP_K_MAX,
                            value=TOP_K_DEFAULT, step=1,
                            label="Top-k"
                        )
                        max_new_tokens = gr.Slider(
                            minimum=MAX_NEW_TOKENS_MIN, maximum=MAX_NEW_TOKENS_MAX,
                            value=MAX_NEW_TOKENS_DEFAULT, step=8,
                            label="Maximum New Tokens"
                        )
        
        # Example prompts
        examples = [
            "Write a short story about a snowflake that comes to life.",
            "Explain the concept of artificial neural networks to a 10-year-old.",
            "What are some interesting applications of natural language processing?",
            "Write a haiku about programming.",
            "Create a dialogue between two AI researchers discussing the future of language models."
        ]

        with gr.Accordion("Example Prompts", open=True):
            with gr.Column(elem_classes="example-section"):
                gr.Examples(
                    examples=examples,
                    inputs=prompt,
                    label="Click on an example to try it",
                    examples_per_page=5
                )
        
        gr.HTML(f"""
        <div class="footer">
            <p>Snowflake Models Demo • Created with Gradio • {datetime.datetime.now().year}</p>
        </div>
        """)
        
        # Interactions
        submit_btn.click(
            fn=generate_text,
            inputs=[prompt, model_version, temperature, top_p, top_k, max_new_tokens, history_state],
            outputs=[response_output, history_state, chat_history_display]
        )
        prompt.submit(
            fn=generate_text,
            inputs=[prompt, model_version, temperature, top_p, top_k, max_new_tokens, history_state],
            outputs=[response_output, history_state, chat_history_display]
        )
        clear_btn.click(
            fn=clear_conversation,
            inputs=[],
            outputs=[prompt, history_state, chat_history_display]
        )

    return demo

# Initialize
print("Loading Snowflake models...")
try:
    load_all_models()
    print("All models loaded successfully!")
    demo = create_demo()
except Exception as e:
    print(f"Error loading models: {e}")
    with gr.Blocks(css=css) as demo:
        gr.HTML(f"""
        <div class="header" style="background-color: #ffebee;">
            <h1><span class="snowflake-icon">⚠️</span> Error Loading Models</h1>
            <p>There was a problem loading the Snowflake models: {str(e)}</p>
        </div>
        """)

# Run app
if __name__ == "__main__":
    demo.launch()