Spaces:

Moklas
/

share

Sleeping

File size: 24,179 Bytes

from flask import Flask, request, render_template_string, Response, stream_with_context
from llama_cpp import Llama
import os
import time
import json

app = Flask(__name__)

# --- ১. সঠিক মডেল লোকেশন ---
# Q1_0 (এন্ড-টু-এন্ড ১-বিট) ফরম্যাট এখন mainline llama.cpp-তে upstream সাপোর্টেড,
# তাই llama-cpp-python (রিসেন্ট ভার্সন) দিয়ে এটা সরাসরি চলে।
# (শুধু তাদের ternary/Q2_0 ফরম্যাটের জন্য বিশেষ ফর্ক লাগে, Q1_0-এর জন্য না।)

# ✅ বর্তমানে সক্রিয়: Bonsai 8B (~1.16 GB)
#MODEL_REPO = "prism-ml/Bonsai-8B-gguf"
#MODEL_FILE = "Bonsai-8B-Q1_0.gguf"

# 🔽 ছোট মডেল চাইলে উপরের দুই লাইন কমেন্ট করে নিচের যেকোনো একটা আনকমেন্ট করুন 🔽

# --- Bonsai 4B (~0.57 GB) — মাঝারি সাইজ, আরও দ্রুত ---
#MODEL_REPO = "prism-ml/Bonsai-4B-gguf"
#MODEL_FILE = "Bonsai-4B-Q1_0.gguf"

# Google gemma3 
MODEL_REPO = "DevQuasar/google.gemma-3n-E2B-it-GGUF"
MODEL_FILE = "google.gemma-3n-E2B-it.Q2_K.gguf"

# --- Bonsai 1.7B (~0.25 GB) — সবচেয়ে ছোট (১ বিলিয়নের কাছাকাছি), সবচেয়ে দ্রুত ---
#MODEL_REPO = "prism-ml/Bonsai-1.7B-gguf"
#MODEL_FILE = "Bonsai-1.7B-Q1_0.gguf"

print("⏳ মডেল ডাউনলোড হচ্ছে... (প্রথমবার একটু সময় লাগবে)")

from huggingface_hub import hf_hub_download
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    local_dir="./models",
    token=None  # পাবলিক রিপো, টোকেন লাগবে না
)
print("✅ ডাউনলোড সম্পূর্ণ!")
print(f"📁 ফাইল লোকেশন: {model_path}")

print("⏳ মডেল লোড হচ্ছে (CPU)...")
try:
    # llama-cpp-python: GGUF মেটাডেটা থেকে নিজেই আর্কিটেকচার (qwen3) বুঝে নেয়,
    # তাই model_type জাতীয় কিছু দিতে হয় না।

    # CPU-এর সব কোর ব্যবহার করা হচ্ছে যাতে prompt processing (prefill) দ্রুত হয়।
    # n_threads     -> টোকেন জেনারেশনের সময় ব্যবহৃত থ্রেড সংখ্যা
    # n_threads_batch -> prompt/prefill প্রসেসিংয়ের সময় ব্যবহৃত থ্রেড সংখ্যা (এটাই প্রথম টোকেনের দেরির মূল কারণ)
    cpu_count = os.cpu_count() or 4

    llm = Llama(
        model_path=model_path,
        n_ctx=2048,
        n_threads=cpu_count,
        n_threads_batch=cpu_count,
        n_batch=512,       # prefill ব্যাচ সাইজ, বড় করলে prompt processing দ্রুত হয়
        verbose=False
    )

    # --- ওয়ার্মআপ ---
    # প্রথম ইনফারেন্স কলে llama.cpp কিছু অভ্যন্তরীণ বাফার/গ্রাফ তৈরি করে যা এক্সট্রা সময় নেয়।
    # সার্ভার চালু হওয়ার সময়ই একটা ডামি জেনারেশন চালিয়ে সেই ওয়ান-টাইম খরচ আগেই সেরে ফেলা হচ্ছে,
    # যাতে ইউজারের প্রথম আসল রিকোয়েস্টে এই পেনাল্টি না লাগে।
    print("🔥 মডেল ওয়ার্মআপ হচ্ছে...")
    list(llm("<|im_start|>user\nহাই<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
              max_tokens=1, stream=True))
    print("✅ ওয়ার্মআপ সম্পন্ন!")

    print("✅ মডেল প্রস্তুত! সার্ভার চালু হচ্ছে...")
except Exception as e:
    print(f"⚠️ llama-cpp-python লোড করতে সমস্যা: {e}")
    llm = None

# --- ২. HTML টেমপ্লেট ---
HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
    <title>Bonsai 8B চ্যাট (Hugging Face Space)</title>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <style>
        * { box-sizing: border-box; }
        body {
            font-family: 'Segoe UI', Tahoma, sans-serif;
            max-width: 900px;
            margin: 20px auto;
            padding: 15px;
            background: #f0f2f5;
        }
        .container {
            background: white;
            border-radius: 16px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            overflow: hidden;
            padding: 20px;
        }
        .header {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 15px 20px;
            margin: -20px -20px 20px -20px;
            border-radius: 16px 16px 0 0;
        }
        .header h2 { margin: 0; font-weight: 400; }
        .header small { opacity: 0.8; font-size: 14px; }

        .chat-box {
            background: #f8f9fa;
            border-radius: 12px;
            padding: 15px;
            height: 400px;
            overflow-y: auto;
            margin-bottom: 15px;
            border: 1px solid #e0e0e0;
        }
        .message {
            margin: 8px 0;
            padding: 10px 15px;
            border-radius: 18px;
            max-width: 80%;
            word-wrap: break-word;
            animation: fadeIn 0.3s ease;
        }
        @keyframes fadeIn {
            from { opacity: 0; transform: translateY(10px); }
            to { opacity: 1; transform: translateY(0); }
        }
        .user-msg {
            background: #007bff;
            color: white;
            margin-left: auto;
            text-align: right;
            border-bottom-right-radius: 4px;
        }
        .bot-msg {
            background: #e9ecef;
            color: #333;
            margin-right: auto;
            border-bottom-left-radius: 4px;
            white-space: pre-wrap;
        }
        .bot-msg.thinking {
            background: #e9ecef;
            color: #666;
            font-style: italic;
        }
        .time-info {
            font-size: 12px;
            color: #4a5568;
            margin: 4px 4px 10px 4px;
            text-align: left;
            background: #f7fafc;
            padding: 6px 12px;
            border-radius: 8px;
            border-left: 3px solid #667eea;
            display: inline-block;
        }
        .time-info strong {
            color: #2d3748;
        }
        .input-area {
            display: flex;
            gap: 10px;
        }
        #userInput {
            flex: 1;
            padding: 12px 18px;
            border: 2px solid #ddd;
            border-radius: 25px;
            font-size: 15px;
            outline: none;
            transition: 0.2s;
        }
        #userInput:focus {
            border-color: #667eea;
        }
        button {
            padding: 12px 28px;
            background: #667eea;
            color: white;
            border: none;
            border-radius: 25px;
            font-size: 15px;
            cursor: pointer;
            transition: 0.2s;
            white-space: nowrap;
        }
        button:hover {
            background: #5a67d8;
            transform: scale(1.02);
        }
        button:disabled {
            background: #a0aec0;
            cursor: not-allowed;
            transform: none;
        }
        .status {
            color: #718096;
            font-size: 13px;
            margin-top: 10px;
            display: flex;
            align-items: center;
            gap: 8px;
        }
        .loader {
            display: inline-block;
            width: 14px;
            height: 14px;
            border: 2px solid #e2e8f0;
            border-top: 2px solid #667eea;
            border-radius: 50%;
            animation: spin 0.8s linear infinite;
        }
        @keyframes spin {
            0% { transform: rotate(0deg); }
            100% { transform: rotate(360deg); }
        }
        .header {
            position: relative;
        }
        .menu-wrapper {
            position: absolute;
            top: 15px;
            right: 20px;
        }
        .menu-btn {
            background: rgba(255,255,255,0.15);
            border: none;
            color: white;
            width: 34px;
            height: 34px;
            padding: 0;
            border-radius: 50%;
            font-size: 20px;
            line-height: 1;
            cursor: pointer;
        }
        .menu-btn:hover {
            background: rgba(255,255,255,0.28);
            transform: none;
        }
        .menu-dropdown {
            display: none;
            position: absolute;
            top: 42px;
            right: 0;
            background: white;
            border-radius: 10px;
            box-shadow: 0 4px 16px rgba(0,0,0,0.2);
            overflow: hidden;
            min-width: 160px;
            z-index: 10;
        }
        .menu-dropdown.open {
            display: block;
        }
        .menu-item {
            padding: 12px 16px;
            color: #333;
            font-size: 14px;
            cursor: pointer;
            white-space: nowrap;
        }
        .menu-item:hover {
            background: #f0f2f5;
        }
        .footer {
            margin-top: 15px;
            text-align: center;
            color: #a0aec0;
            font-size: 12px;
        }
        .first-token-badge {
            background: #48bb78;
            color: white;
            padding: 2px 10px;
            border-radius: 12px;
            font-size: 11px;
            margin-left: 8px;
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h2>🌳 Bonsai 8B (১-বিট) চ্যাট</h2>
            <small>Hugging Face Space • CPU Inference • মডেল সাইজ: ~১.১৫ GB (Q1_0, 1-bit)</small>
            <div class="menu-wrapper">
                <button id="menuBtn" class="menu-btn">⋮</button>
                <div id="menuDropdown" class="menu-dropdown">
                    <div id="newChatItem" class="menu-item">➕ নতুন চ্যাট</div>
                </div>
            </div>
        </div>

        <div class="chat-box" id="chatBox">
            <div class="message bot-msg">👋 হ্যালো! আমি Bonsai 8B। আপনি কী জানতে চান?</div>
        </div>

        <div class="input-area">
            <input type="text" id="userInput" placeholder="এখানে প্রশ্ন লিখুন..." />
            <button id="sendBtn">পাঠান</button>
        </div>
        <div class="status" id="status">
            <span>✅</span> প্রস্তুত
        </div>
        <div class="footer">
            Powered by llama-cpp-python • Bonsai 8B (Q1_0, 1-bit)
        </div>
    </div>

    <script>
        const chatBox = document.getElementById('chatBox');
        const input = document.getElementById('userInput');
        const sendBtn = document.getElementById('sendBtn');
        const status = document.getElementById('status');
        const menuBtn = document.getElementById('menuBtn');
        const menuDropdown = document.getElementById('menuDropdown');
        const newChatItem = document.getElementById('newChatItem');

        // সর্বশেষ ৩টা প্রশ্ন-উত্তরের হিস্টরি রাখা হচ্ছে (মডেলকে কনটেক্সট দেওয়ার জন্য)
        let conversationHistory = [];

        function addMessage(text, isUser = false, isThinking = false) {
            const div = document.createElement('div');
            div.className = `message ${isUser ? 'user-msg' : 'bot-msg'}${isThinking ? ' thinking' : ''}`;
            div.textContent = text;
            chatBox.appendChild(div);
            chatBox.scrollTop = chatBox.scrollHeight;
            return div;
        }

        function addTimeInfo(firstTokenTime, totalTime, startTimeStr, endTimeStr) {
            const div = document.createElement('div');
            div.className = 'time-info';
            div.innerHTML = `
                <strong>⏱️ প্রথম টোকেন:</strong> ${firstTokenTime} সেকেন্ড &nbsp;|&nbsp;
                <strong>মোট সময়:</strong> ${totalTime} সেকেন্ড &nbsp;|&nbsp;
                <strong>শুরু:</strong> ${startTimeStr} &nbsp;|&nbsp;
                <strong>শেষ:</strong> ${endTimeStr}
            `;
            chatBox.appendChild(div);
            chatBox.scrollTop = chatBox.scrollHeight;
        }

        async function sendMessage() {
            const text = input.value.trim();
            if (!text) return;

            // সেন্ড বাটন ক্লিক করার মুহূর্ত থেকেই টাইমার শুরু
            const startTime = Date.now();

            addMessage(text, true);
            input.value = '';
            sendBtn.disabled = true;
            status.innerHTML = '<span class="loader"></span> প্রক্রিয়াকরণ শুরু হচ্ছে...';

            // থিংকিং ইন্ডিকেটর
            const thinkingDiv = addMessage('⏳ উত্তর তৈরি হচ্ছে...', false, true);

            try {
                const response = await fetch('/chat', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ prompt: text, history: conversationHistory })
                });

                if (!response.ok || !response.body) {
                    throw new Error('স্ট্রিম শুরু করা যায়নি');
                }

                // থিংকিং ইন্ডিকেটর সরানো
                thinkingDiv.remove();
                const botDiv = addMessage('', false);

                const reader = response.body.getReader();
                const decoder = new TextDecoder();
                let fullText = '';
                let buffer = '';
                let metaFound = false;
                let firstToken = true;
                let firstTokenTime = 0;

                while (true) {
                    const { done, value } = await reader.read();
                    if (done) break;
                    buffer += decoder.decode(value, { stream: true });

                    const metaIndex = buffer.indexOf('[[META]]');
                    if (metaIndex !== -1) {
                        fullText += buffer.slice(0, metaIndex);
                        botDiv.textContent = fullText;
                        buffer = buffer.slice(metaIndex);
                        metaFound = true;
                    } else if (!metaFound) {
                        fullText += buffer;
                        botDiv.textContent = fullText;
                        buffer = '';
                        
                        // প্রথম টোকেন আসার সময় রেকর্ড
                        if (firstToken && fullText.length > 0) {
                            firstToken = false;
                            firstTokenTime = (Date.now() - startTime) / 1000;
                            status.innerHTML = `✍️ উত্তর লেখা হচ্ছে... <span class="first-token-badge">প্রথম টোকেন: ${firstTokenTime.toFixed(2)}s</span>`;
                        }
                    }
                    chatBox.scrollTop = chatBox.scrollHeight;
                }

                // মেটাডেটা পার্স করা
                const metaIndex = buffer.indexOf('[[META]]');
                if (metaIndex !== -1) {
                    try {
                        const meta = JSON.parse(buffer.slice(metaIndex + 8));
                        // firstTokenTime যদি সেট না হয়ে থাকে (খুব ছোট রেসপন্স)
                        if (firstTokenTime === 0) {
                            firstTokenTime = meta.first_token_time || 0;
                        }
                        addTimeInfo(
                            firstTokenTime.toFixed(2),
                            meta.elapsed.toFixed(2),
                            meta.start,
                            meta.end
                        );
                    } catch (e) {
                        // মেটাডেটা পার্স করতে ব্যর্থ হলে চুপচাপ বাদ দেওয়া হলো
                    }
                }

                chatBox.scrollTop = chatBox.scrollHeight;
                status.innerHTML = '✅ প্রস্তুত';

                // এই এক্সচেঞ্জটা হিস্টরিতে যোগ করা এবং শুধু সর্বশেষ ৩টা রাখা
                conversationHistory.push({ user: text, assistant: fullText.trim() });
                if (conversationHistory.length > 3) {
                    conversationHistory = conversationHistory.slice(-3);
                }
            } catch (error) {
                thinkingDiv.remove();
                addMessage('⚠️ সার্ভার ত্রুটি! আবার চেষ্টা করুন।', false);
                status.innerHTML = '❌ সংযোগ ত্রুটি';
                console.error('Error:', error);
            }
            sendBtn.disabled = false;
            input.focus();
        }

        function newChat() {
            chatBox.innerHTML = '';
            conversationHistory = [];
            addMessage('👋 হ্যালো! আমি Bonsai 8B। আপনি কী জানতে চান?', false);
            status.innerHTML = '✅ প্রস্তুত';
            menuDropdown.classList.remove('open');
            input.focus();
        }

        sendBtn.addEventListener('click', sendMessage);
        input.addEventListener('keypress', (e) => {
            if (e.key === 'Enter') sendMessage();
        });

        // থ্রি-ডট মেনু টগল
        menuBtn.addEventListener('click', (e) => {
            e.stopPropagation();
            menuDropdown.classList.toggle('open');
        });
        newChatItem.addEventListener('click', newChat);
        // মেনুর বাইরে ক্লিক করলে বন্ধ হয়ে যাবে
        document.addEventListener('click', () => {
            menuDropdown.classList.remove('open');
        });

        // Auto-focus on page load
        input.focus();
    </script>
</body>
</html>
"""

# --- ৩. Flask রাউট ---
@app.route('/')
def index():
    return render_template_string(HTML_TEMPLATE)

@app.route('/chat', methods=['POST'])
def chat():
    data = request.get_json()
    user_prompt = data.get('prompt', '')
    # ফ্রন্টএন্ড থেকে পাঠানো সর্বশেষ (সর্বোচ্চ ৩টা) প্রশ্ন-উত্তরের হিস্টরি
    # প্রতিটা আইটেম: {"user": "...", "assistant": "..."}
    history = data.get('history', [])

    if not user_prompt:
        return {"error": "কোনো প্রশ্ন নেই"}

    if llm is None:
        return {"error": "মডেল লোড হয়নি। সার্ভার লগ চেক করুন।"}

    # শুধু সর্বশেষ ৩টা এক্সচেঞ্জ ব্যবহার করা হচ্ছে (কনটেক্সট উইন্ডো সীমিত রাখতে)
    if isinstance(history, list):
        history = history[-3:]
    else:
        history = []

    # থিংকিং সম্পূর্ণভাবে বন্ধ — খালি <think></think> প্রি-ফিল করে মডেলকে সরাসরি উত্তরে পাঠানো হচ্ছে
    # হিস্টরির প্রতিটা টার্ন প্রম্পটে জুড়ে দেওয়া হচ্ছে, যাতে মডেল আগের কনভারসেশন মনে রাখে
    full_prompt = ""
    for turn in history:
        h_user = str(turn.get('user', '')).strip()
        h_assistant = str(turn.get('assistant', '')).strip()
        if not h_user:
            continue
        full_prompt += f"<|im_start|>user\n{h_user}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n{h_assistant}<|im_end|>\n"

    full_prompt += f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"

    def generate():
        start_time = time.time()
        start_str = time.strftime('%H:%M:%S', time.localtime(start_time))
        first_token_time = None
        
        try:
            # stream=True দিলে llama-cpp-python টোকেন-বাই-টোকেন জেনারেট করে দেয়
            stream = llm(
                full_prompt,
                max_tokens=256,          # সর্বোচ্চ টোকেন সংখ্যা
                temperature=0.1,         # তাপমাত্রা (কম = ডিটারমিনিস্টিক)
                top_p=0.2,              # নিউক্লিয়াস স্যাম্পলিং
                top_k=20,               # টপ-কে স্যাম্পলিং
                repeat_penalty=1.4,     # পুনরাবৃত্তি শাস্তি
                stop=["<|im_end|>", "user:", "User:"],
                stream=True
            )
            
            for chunk in stream:
                token_text = chunk["choices"][0]["text"]
                if token_text:
                    # প্রথম টোকেনের সময় রেকর্ড
                    if first_token_time is None:
                        first_token_time = time.time() - start_time
                    yield token_text
        except Exception as e:
            print(f"❌ জেনারেশন ত্রুটি: {e}")
            yield f"\n⚠️ মডেল ত্রুটি: {str(e)}"
        finally:
            end_time = time.time()
            end_str = time.strftime('%H:%M:%S', time.localtime(end_time))
            elapsed = round(end_time - start_time, 2)
            first_token = round(first_token_time, 2) if first_token_time else 0
            
            print(f"⏱️ প্রথম টোকেন: {first_token} সেকেন্ড | মোট সময়: {elapsed} সেকেন্ড")
            print(f"📊 প্যারামিটার: max_tokens=256, temp=0.1, top_p=0.2, top_k=40, repeat_penalty=1.1")
            
            # স্ট্রিমের একদম শেষে টাইমিং তথ্য পাঠানো
            meta = {
                "start": start_str,
                "end": end_str,
                "elapsed": elapsed,
                "first_token_time": first_token
            }
            yield f"\n[[META]]{json.dumps(meta, ensure_ascii=False)}"

    return Response(stream_with_context(generate()), mimetype='text/plain')

# --- ৪. সার্ভার চালানো ---
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)