app.py

import spaces
import gradio as gr
import torch
from datetime import datetime
import os
import subprocess # For Flash Attention install
from threading import Thread # For streaming

# --- Install Flash Attention (specific method for compatibility) ---
print("Attempting to install Flash Attention 2...")
try:
    subprocess.run(
        'pip install flash-attn --no-build-isolation',
        env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
        shell=True,
        check=True
    )
    print("Flash Attention installed successfully using subprocess method.")
    _flash_attn_2_available = True
except Exception as e:
    print(f"Could not install Flash Attention 2 using subprocess: {e}")
    print("Proceeding without Flash Attention 2. Performance may be impacted.")
    _flash_attn_2_available = False

# --- Import Transformers AFTER potential install ---
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer # Added TextIteratorStreamer
from huggingface_hub import HfApi, HfFolder

# --- Configuration ---
model_id = "Tesslate/UIGEN-T3-4B-Preview"
creator_link = "https://huggingface.co/TesslateAI"
model_link = f"https://huggingface.co/{model_id}"
website_link = "https://tesslate.com"
discord_link = "https://discord.gg/DkzMzwBTaw"

# --- Text Content (Keep the cool UI elements) ---
Title = f"""
<div style="text-align: center; margin-bottom: 20px;">
    <img src="https://huggingface.co/Tesslate/Tessa-T1-14B/resolve/main/tesslate_logo_color.png?download=true" alt="Tesslate Logo" style="height: 80px; margin-bottom: 10px;">
    <h1 style="margin-bottom: 5px;">🚀 Welcome to the UIGEN Playground 🚀</h1>
    <p style="font-size: 1.1em;">Specialized UI based reasoning</p>
    <p>Model by <a href="{creator_link}" target="_blank">TesslateAI</a> | <a href="{model_link}" target="_blank">View on Hugging Face</a> | Running with 8-bit Quantization | Streaming Output</p>
</div>
"""

description = f"""
Interact with **[{model_id}]({model_link})**.
UIGEN-T2 specializes in **frontend development**, leveraging advanced reasoning to autonomously generate well-structured components.
This demo uses **8-bit quantization** via `bitsandbytes` for reduced memory footprint. **Flash Attention 2** is enabled if available. Output is **streamed** token-by-token.
"""

# --- (Keep about_tesslate and join_us sections as before) ---
about_tesslate = f"""
## About Tesslate & Our Vision
<img src="https://huggingface.co/Tesslate/Tessa-T1-14B/resolve/main/tesslate_logo_notext.png?download=true" alt="Tesslate Icon" style="height: 40px; float: left; margin-right: 10px;">
Hi everyone, I’m Manav, founder of Tesslate, and we’re on a mission to revolutionize AI by putting powerful reasoning models into your hands.

Today, the AI landscape is dominated by massive frontier models—large, costly, and slow. At Tesslate, we see things differently. The next wave of AI disruption won’t come from sheer size; it'll be driven by **speed, specialization, and precision reasoning**. Smaller, specialized models aren’t just faster—they’re smarter and more efficient.

Our story began when we released a UI-generation model on Hugging Face that didn't just replicate patterns—it could reason through entire component hierarchies. It resonated instantly, hitting over 10,000 downloads in weeks. That early success validated our vision, and we doubled down.

At Tesslate, we build lean, intelligent models that:
*   🧠 **Think** like human agents
*   💡 **Reason** through complex, real-world workflows
*   💻 **Execute** like elite developers, designers, and analysts

We've already delivered:
*   **UIGEN-T1.5:** Creating stunning, editable interfaces (React, Tailwind, Three.js)
*   **Tessa-T1:** A specialized reasoning engine optimized for React development and AI agents (You are here!)
*   **Synthia S1:** Our flagship general-reasoning model, proving powerful reasoning capabilities beyond STEM into creativity and storytelling.

Our vision is bigger. We aim to do everything covering training, inference, real-time agent actions, infrastructure, research, and innovative products. We’re already piloting with industry-leading clients tackling everything from sophisticated design systems to real-time analytics.

**Join us!** We're seeking strategic advice, introductions, compute resources, and capital.
👉 Visit **[tesslate.com]({website_link})** to learn more and connect.
"""

join_us = f"""
<div style="text-align: center;">
    <h3 style="margin-bottom: 10px;">Connect with Tesslate</h3>
    <a href="{discord_link}" target="_blank" style="text-decoration: none; margin: 0 10px;">
        <img src="https://img.shields.io/discord/1225631184402124842?label=Discord&logo=discord&style=for-the-badge&color=5865F2" alt="Join us on Discord">
    </a>
    <a href="{website_link}" target="_blank" style="text-decoration: none; margin: 0 10px;">
        <img src="https://img.shields.io/badge/Website-tesslate.com-blue?style=for-the-badge&logo=googlechrome&logoColor=white" alt="Visit tesslate.com">
    </a>
     <a href="{model_link}" target="_blank" style="text-decoration: none; margin: 0 10px;">
        <img src="https://img.shields.io/badge/🤗%20Model-Tessa--T1--14B-yellow?style=for-the-badge&logo=huggingface" alt="Tessa-T1-14B on Hugging Face">
    </a>
</div>
"""
# --- Model and Tokenizer Loading ---
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
if device == "cpu":
    print("Warning: Running on CPU. Quantization and Flash Attention require CUDA.")
    _flash_attn_2_available = False

hf_token = os.getenv('HF_TOKEN')
if not hf_token:
    try:
        hf_token = HfFolder.get_token()
        if not hf_token: hf_token = HfApi().token
        if not hf_token: raise ValueError("HF token not found.")
        print("Using token from Hugging Face login.")
    except Exception as e:
         raise ValueError(f"HF token acquisition failed: {e}. Please set HF_TOKEN or login.")

print(f"Loading Tokenizer: {model_id}")
tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token, trust_remote_code=True)

print(f"Loading Model: {model_id} with 8-bit quantization")
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
attn_implementation = "flash_attention_2" if _flash_attn_2_available and device == "cuda" else "sdpa"
print(f"Using attention implementation: {attn_implementation}")

try:
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        token=hf_token,
        device_map="auto",
        quantization_config=quantization_config,
        attn_implementation=attn_implementation,
        trust_remote_code=True
    )
    print("Model loaded successfully with 8-bit quantization.")
except Exception as e:
    print(f"Error loading model: {e}")
    if attn_implementation == "flash_attention_2":
        print("Flash Attention 2 failed at load time. Trying fallback 'sdpa' attention...")
        try:
            attn_implementation = "sdpa"
            model = AutoModelForCausalLM.from_pretrained(
                model_id, token=hf_token, device_map="auto", quantization_config=quantization_config,
                attn_implementation=attn_implementation, trust_remote_code=True
            )
            print("Model loaded successfully with 8-bit quantization and SDPA attention.")
        except Exception as e2:
             print(f"Fallback to SDPA attention also failed: {e2}"); raise e2
    else: raise e

# --- (Keep config info gathering and tokenizer info formatting as before) ---
try:
    config_json = model.config.to_dict()
    quant_info = model.config.quantization_config.to_dict() if hasattr(model.config, 'quantization_config') else {}
    model_config_info = f"""
**Model Type:** {config_json.get('model_type', 'N/A')}
**Architecture:** {config_json.get('architectures', ['N/A'])[0]}
**Vocab Size:** {config_json.get('vocab_size', 'N/A')}
**Hidden Size:** {config_json.get('hidden_size', 'N/A')}
**Num Hidden Layers:** {config_json.get('num_hidden_layers', 'N/A')}
**Num Attention Heads:** {config_json.get('num_attention_heads', 'N/A')}
**Max Position Embeddings:** {config_json.get('max_position_embeddings', 'N/A')}
**Attention Implementation:** `{attn_implementation}`
**Quantization:** 8-bit (`load_in_8bit={quant_info.get('load_in_8bit', 'N/A')}`)
"""
except Exception as e:
    print(f"Could not retrieve full model config: {e}")
    model_config_info = f"**Error:** Could not load full config details for {model_id}."

def format_tokenizer_info(tokenizer_instance):
    try:
        info = [
            f"**Tokenizer Class:** `{tokenizer_instance.__class__.__name__}`",
            f"**Vocabulary Size:** {tokenizer_instance.vocab_size}",
            f"**Model Max Length:** {tokenizer_instance.model_max_length}",
            f"**EOS Token:** `{tokenizer_instance.eos_token}` (ID: {tokenizer_instance.eos_token_id})",
            f"**Special Tokens:** Check model card for specific template/tokens.", # Qwen2 has specific tokens
        ]
        # Add BOS/PAD/UNK if they are commonly used and different from EOS
        if hasattr(tokenizer_instance, 'pad_token') and tokenizer_instance.pad_token and tokenizer_instance.pad_token_id is not None:
             info.append(f"**Padding Token:** `{tokenizer_instance.pad_token}` (ID: {tokenizer_instance.pad_token_id})")
        if hasattr(tokenizer_instance, 'bos_token') and tokenizer_instance.bos_token and tokenizer_instance.bos_token_id is not None:
             info.append(f"**BOS Token:** `{tokenizer_instance.bos_token}` (ID: {tokenizer_instance.bos_token_id})")
        if hasattr(tokenizer_instance, 'unk_token') and tokenizer_instance.unk_token and tokenizer_instance.unk_token_id is not None:
             info.append(f"**UNK Token:** `{tokenizer_instance.unk_token}` (ID: {tokenizer_instance.unk_token_id})")
        return "\n".join(info)
    except Exception as e:
        print(f"Error getting tokenizer info: {e}")
        return f"Could not retrieve full tokenizer details. Vocab size: {getattr(tokenizer_instance, 'vocab_size', 'N/A')}"

tokenizer_info = format_tokenizer_info(tokenizer)


# --- Generation Function (Modified for Streaming) ---
@spaces.GPU(duration=600)
def generate_response(system_prompt, user_prompt, temperature, max_new_tokens, top_p, repetition_penalty, top_k, min_p):
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": user_prompt})

    try:
        full_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    except Exception as e:
        print(f"Warning: Using fallback prompt format due to error: {e}")
        prompt_parts = []
        if system_prompt and system_prompt.strip(): prompt_parts.append(f"System: {system_prompt}")
        prompt_parts.append(f"\nUser: {user_prompt}\nAssistant:")
        full_prompt = "\n".join(prompt_parts)

    # Use TextIteratorStreamer for streaming output
    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=10.0, # Timeout for waiting for new tokens
        skip_prompt=True, # Don't yield the prompt
        skip_special_tokens=True
    )

    # Ensure inputs are correctly placed (device_map handles this)
    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=4096).to(model.device) # Use model's device

    # Generation kwargs, pass streamer
    generation_kwargs = dict(
        inputs, # Pass tokenized inputs directly
        streamer=streamer, # Pass the streamer
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature) if float(temperature) > 0 else None,
        top_p=float(top_p),
        top_k=int(top_k),
        repetition_penalty=float(repetition_penalty),
        do_sample=True if float(temperature) > 0 else False,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

    if temperature == 0:
        generation_kwargs.pop('top_p', None)
        generation_kwargs.pop('top_k', None)
        generation_kwargs['do_sample'] = False

    # Run generation in a separate thread
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # Yield generated text as it becomes available
    generated_text = ""
    # Yield an empty string immediately to clear previous output
    yield ""
    for new_text in streamer:
        generated_text += new_text
        yield generated_text

# --- Gradio Interface (No changes needed here for streaming itself) ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), css=".gradio-container { max-width: 90% !important; }") as demo:
    gr.Markdown(Title)
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Group():
                system_prompt = gr.Textbox(
                    label="System Prompt (Persona & Instructions)",
                    value="You are Tesslate, a helpful assistant specialized in UI generation.",
                    lines=3,
                    info="Guide the model's overall behavior and expertise."
                )
                user_prompt = gr.Textbox(
                    label="💬 Your Request",
                    placeholder="e.g., 'Create a dashboard for my dog washing business' or 'Show me an airbnb clone'",
                    lines=6
                )

            with gr.Accordion("🛠️ Generation Parameters", open=True):
                 with gr.Row():
                     temperature = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.05, label="🌡️ Temperature")
                     max_new_tokens = gr.Slider(minimum=64, maximum=10000, value=10000, step=32, label="📊 Max New Tokens")
                 with gr.Row():
                     top_k = gr.Slider(minimum=1, maximum=200, value=40, step=1, label="🏆 Top-k")
                     top_p = gr.Slider(minimum=0.05, maximum=1.0, value=0.95, step=0.01, label="🏅 Top-p (nucleus)")
                 with gr.Row():
                     repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, value=1.1, step=0.01, label="🦜 Repetition Penalty")
                     min_p = gr.Slider(minimum=0.0, maximum=0.5, value=0.05, step=0.01, label="📉 Min-p (Not Active)")

            generate_btn = gr.Button("🚀 Generate Response (Streaming)", variant="primary", size="lg") # Updated button text slightly

        with gr.Column(scale=2):
            output = gr.Code(
                label=f"🌠 Tessa-T1-14B (8-bit) Output",
                language="markdown",
                lines=25,
                # interactive=False # Usually keep interactive=False for Code output
            )

            with gr.Accordion("⚙️ Model & Tokenizer Details", open=False):
                 gr.Markdown("### Model Configuration")
                 gr.Markdown(model_config_info)
                 gr.Markdown("---")
                 gr.Markdown("### Tokenizer Configuration")
                 gr.Markdown(tokenizer_info)

    # --- (Keep About Tesslate, Links, and Examples sections as before) ---
    with gr.Row():
        with gr.Accordion("💡 About Tesslate & Our Mission", open=False):
            gr.Markdown(about_tesslate)

    gr.Markdown(join_us)

    gr.Examples(
        examples=[
            [
                "You are Tesslate, a helpful assistant specialized in UI generation.",
                "Make a really good looking dashboard with charts.",
                0.7, 512, 0.95, 1.1, 40, 0.05
            ],
            [
                "You are Tesslate, a helpful assistant specialized in UI generation.",
                "Make an animated navbar.",
                0.7, 1024, 0.95, 1.1, 40, 0.05
            ],
            [
                "You are Tesslate, a helpful assistant specialized in UI generation.",
                "Make an Airbnb clone.",
                0.7, 1536, 0.95, 1.1, 40, 0.05
            ],
             [
                "You are Tesslate, a helpful assistant specialized in UI generation.",
                "Create a special website.",
                0.8, 1024, 0.98, 1.05, 60, 0.05
            ]
        ],
        inputs=[
            system_prompt,
            user_prompt,
            temperature,
            max_new_tokens,
            top_p,
            repetition_penalty,
            top_k,
            min_p
        ],
        outputs=output,
        label="✨ Example Prompts (Click to Load)"
    )

    # --- Connect button click to the GENERATOR function ---
    generate_btn.click(
        fn=generate_response,
        inputs=[system_prompt, user_prompt, temperature, max_new_tokens, top_p, repetition_penalty, top_k, min_p],
        outputs=output,
        api_name="generate_stream" # Changed API name for clarity
    )

# --- Launch the demo ---
if __name__ == "__main__":
    demo.queue().launch(debug=True, share=False)