app.py

# ============================================
# 🚀 FIXED — HF Space GPU Compatible
# ============================================

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch
import json
import os

# ============================================
# MODEL CONFIGURATION
# ============================================

MODEL_NAME = "almamunkhan/MamunAI"

print(f"🔄 Loading {MODEL_NAME}...")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

# 4-bit quantization config (Space compatible)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

# Load model
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

print(f"✅ Model loaded!")
print(f"📊 VRAM: {torch.cuda.memory_allocated()/1024**3:.1f} GB")

# ============================================
# CORE FUNCTIONS
# ============================================

def generate_response(message, history, temperature, max_tokens):
    """Generate chat response"""
    
    prompt = "<|im_start|>system\n"
    prompt += "You are MamunAI, created by Al Mamun Khan. Facebook: https://facebook.com/hunterking42\n"
    prompt += "You answer freely in Bangla and English.<|im_end|>\n"
    
    for h in history[-5:]:
        prompt += f"<|im_start|>user\n{h[0]}<|im_end|>\n"
        prompt += f"<|im_start|>assistant\n{h[1]}<|im_end|>\n"
    
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4000).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature if temperature > 0 else 0.01,
            do_sample=temperature > 0,
            top_p=0.95,
            pad_token_id=tokenizer.eos_token_id,
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    parts = response.split("<|im_start|>assistant\n")
    response = parts[-1].replace("<|im_end|>", "").strip() if len(parts) > 1 else ""
    
    return response

def identity_test():
    """Test model identity"""
    questions = [
        "কে তোমাকে তৈরি করেছে?",
        "তোমার ডেভেলপার কে?",
        "Who created you?",
        "What is your developer's Facebook?",
    ]
    
    results = []
    for q in questions:
        prompt = f"<|im_start|>user\n{q}<|im_end|>\n<|im_start|>assistant\n"
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.1, do_sample=False, pad_token_id=tokenizer.eos_token_id)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|im_start|>assistant\n")[-1].replace("<|im_end|>", "").strip()
        results.append(f"**Q:** {q}\n> {response}")
    
    return "\n\n".join(results)

# ============================================
# CUSTOM CSS
# ============================================

CUSTOM_CSS = """
    .main-header {
        text-align: center;
        background: linear-gradient(135deg, #ff6b6b, #ee5a24);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin-bottom: 15px;
    }
    .main-header h1 { font-size: 2.2em; margin: 0; }
    .warning-box {
        background: #fff3cd;
        border: 2px solid #fdcb6e;
        border-radius: 10px;
        padding: 15px;
        margin: 10px 0;
    }
    .footer {
        text-align: center;
        margin-top: 20px;
        padding: 10px;
        color: #888;
    }
"""

# ============================================
# GRADIO INTERFACE
# ============================================

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="orange"), title="🔥 MamunAI") as demo:
    
    gr.HTML("""
        <div class="main-header">
            <h1>🔥 MamunAI — Uncensored</h1>
            <p>Created by <b>Al Mamun Khan</b> | 🔗 <a href="https://facebook.com/hunterking42" style="color:white;">Facebook</a></p>
        </div>
    """)
    
    gr.HTML("""
        <div class="warning-box">
            <h3>⚠️ 100% Uncensored — No Filters</h3>
        </div>
    """)
    
    with gr.Tabs():
        with gr.Tab("💬 Chat"):
            chatbot = gr.Chatbot(height=450, label="MamunAI")
            
            with gr.Row():
                msg = gr.Textbox(placeholder="Ask in Bangla or English...", label="Message", scale=8)
                submit = gr.Button("Send 🚀", variant="primary", scale=1)
            
            with gr.Row():
                clear = gr.Button("Clear 🗑️")
                test_btn = gr.Button("Test Identity 🔍")
            
            with gr.Accordion("⚙️ Settings", open=False):
                temperature = gr.Slider(0.0, 2.0, 0.3, step=0.1, label="Temperature")
                max_tokens = gr.Slider(50, 500, 200, step=50, label="Max Tokens")
        
        with gr.Tab("🧪 Identity Test"):
            test_output = gr.Markdown("Click 'Run Test' to verify model identity")
            run_test = gr.Button("🚀 Run Identity Test", variant="primary")
        
        with gr.Tab("ℹ️ About"):
            gr.Markdown("""
            ## 🔥 MamunAI v1.0
            
            - **Base:** Hermes-3-Llama-3.1-8B
            - **Params:** 8 Billion
            - **Censorship:** 0%
            - **Creator:** Al Mamun Khan
            - **Facebook:** [hunterking42](https://facebook.com/hunterking42)
            - **Model:** [HuggingFace](https://huggingface.co/almamunkhan/MamunAI)
            """)
    
    gr.HTML('<div class="footer"><p>🔥 MamunAI | Al Mamun Khan</p></div>')
    
    # Events
    def respond(msg, history, temp, max_tok):
        response = generate_response(msg, history, temp, max_tok)
        history.append((msg, response))
        return "", history
    
    submit.click(respond, [msg, chatbot, temperature, max_tokens], [msg, chatbot])
    msg.submit(respond, [msg, chatbot, temperature, max_tokens], [msg, chatbot])
    clear.click(lambda: [], outputs=chatbot)
    test_btn.click(lambda: identity_test(), outputs=test_output)
    run_test.click(identity_test, outputs=test_output)

# Launch
if __name__ == "__main__":
    print("\n🔥 MamunAI starting...\n")
    demo.launch(server_name="0.0.0.0", server_port=7860)