app.py

import gradio as gr
import subprocess
import os
import tempfile
from gtts import gTTS
from transformers import pipeline

# Ensure execution context is inside the compiled architecture directory
if os.path.exists("/app"):
    os.chdir("/app")
elif os.path.exists("/content/BitNet"):
    os.chdir("/content/BitNet")

# ==============================================================================
# INITIALIZE ASR ENGINE
# ==============================================================================
print("[SYSTEM] Loading Whisper ASR Engine...")
asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=-1)

# ==============================================================================
# CONSTANTS & CONFIGURATION
# ==============================================================================
MODEL_PATH = "models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf"

DEFAULT_SYSTEM_PROMPT = (
    "You are a Socratic assistant. Do not answer questions directly. "
    "Instead, respond exclusively with 3 deep, reflective questions. "
    "Then generate %^%^%^"
)

# Advanced Voice Matrix routing for Edge-TTS
VOICE_MATRIX = {
    "English": {"Male": "en-US-ChristopherNeural", "Female": "en-US-AvaNeural", "gtts": "en"},
    "German": {"Male": "de-DE-ChristophNeural", "Female": "de-DE-KatjaNeural", "gtts": "de"}
}

# ==============================================================================
# TTS AUDIO GENERATOR
# ==============================================================================
def generate_tts_audio(text, lang="English", gender="Male", engine="edge-tts"):
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    temp_path = temp_file.name
    temp_file.close()
    
    try:
        if engine == "edge-tts":
            voice = VOICE_MATRIX[lang][gender]
            result = subprocess.run([
                "edge-tts", 
                "--voice", voice, 
                "--text", text, 
                "--write-media", temp_path
            ], capture_output=True, text=True)
            if result.returncode != 0:
                print(f"[TTS ERROR] {result.stderr}")
                return None
        else:
            tts = gTTS(text=text, lang=VOICE_MATRIX[lang]["gtts"], slow=False)
            tts.save(temp_path)
            
        if os.path.exists(temp_path) and os.path.getsize(temp_path) > 0:
            return temp_path
        return None
    except Exception as e:
        print(f"[TTS FATAL ERROR] {str(e)}")
        return None

# ==============================================================================
# STREAMING ENGINE
# ==============================================================================
def streaming_chat(text_input, audio_input, lang, gender, system_prompt, tts_engine):
    # Determine input source (Audio takes priority if both exist)
    user_query = text_input
    if audio_input is not None:
        yield "🎙️ Transcribing audio...", gr.skip()
        asr_res = asr_pipe(audio_input, generate_kwargs={"language": lang.lower()})
        user_query = asr_res["text"]
        
    if not user_query or not user_query.strip():
        yield "Please provide text or audio input.", gr.skip()
        return

    formatted_chat_prompt = f"System: {system_prompt}\nUser: {user_query}\nAssistant:"
    
    cmd = [
        "python3", "run_inference.py",
        "-m", MODEL_PATH,
        "-p", formatted_chat_prompt,
        "-n", "120",
        "-temp", "0.4",
        "-t", "2"
    ]
    
    process = subprocess.Popen(
        cmd, 
        stdout=subprocess.PIPE, 
        stderr=subprocess.DEVNULL,
        text=True, 
        bufsize=1
    )
    
    accumulator = ""
    prompt_cleared = False
    LOOKAHEAD_SIZE = 45 
    stop_markers = ["%^%^%^","%%%","[end of text]", "User:", "Assistant:"]
    
    while True:
        char = process.stdout.read(1)
        if not char:
            break
            
        accumulator += char
        
        if not prompt_cleared:
            if "Assistant:" in accumulator:
                prompt_cleared = True
                accumulator = accumulator.split("Assistant:")[-1].lstrip()
            continue
            
        stop_triggered = False
        for marker in stop_markers:
            if marker in accumulator:
                accumulator = accumulator.split(marker)[0]
                stop_triggered = True
                break
                
        if stop_triggered:
            process.terminate()
            break
            
        if len(accumulator) > LOOKAHEAD_SIZE:
            safe_to_display = accumulator[:len(accumulator) - LOOKAHEAD_SIZE]
            yield safe_to_display.strip(), gr.skip()
            
    final_text = accumulator.strip()
    
    if final_text:
        audio_path = generate_tts_audio(final_text, lang, gender, engine=tts_engine)
        yield final_text, audio_path
    else:
        yield final_text, gr.skip()

# ==============================================================================
# TECHNICAL REPORT MARKDOWN TEXT
# ==============================================================================
TECHNICAL_REPORT_MD = """
## 📋 Technical Report: 1-Bit LLM Socratic Refinement Pipeline
**Architecture Core:** Ternary Quantized (1.58-bit) Matrix Processing  

---

### 1. Executive Objective & Target Dataset
The goal of this initiative was to engineer a hyper-lightweight, lightning-fast edge computing application capable of engaging users in conversational Socratic exploration. Traditional full-precision models require significant memory overhead to hold nuanced philosophical frameworks. This project focused on building an ultra-compressed conversational experience capable of executing inside a constrained local CPU footprint (e.g., standard consumer laptops or free cloud application tiers).

* **Target Fine-Tuning Dataset:** `sanjaypantdsd/socratic-method-conversations`
* **Data Characteristics:** High-quality, clean input-to-output mappings that translate explicit factual questions or structural concepts directly into clusters of exactly three open-ended, deeply analytical questions.

---

### 2. Model Training Matrix & Evaluation Phase
Our initial strategy focused on fine-tuning custom models directly on our targeted Socratic dataset. The results exposed clear engineering trade-offs:

| Model Identifier | Architecture Configuration | Operational Performance | Qualitative Evaluation |
| :--- | :--- | :--- | :--- |
| **st192011/bitnet-socratic-1.58b** | Full-precision parameter adjustments tailored to target dataset. | **Excellent** | Produced highly coherent Socratic question arrays aligning perfectly with training structures. |
| **st192011/socratic-bitnet-2b** | Quantized Ternary Representation Variant of custom weights. | **Critically Poor** | Suffered extreme degradation. The model experienced severe structural collapse, outputting infinite semantic loops or unreadable token gibberish. |

#### Analysis of Quantization Collapse
The stark failure of `st192011/socratic-bitnet-2b` highlights a common hurdle in customized 1-bit AI development. When a model's weights are aggressively compressed down to simple ternary values (-1, 0, 1), the mathematical boundaries become extremely rigid. Standard quantization tools often distort the delicate behavioral traits introduced during fine-tuning.

---

### 3. Strategy Pivot: Pretrained Weights + Structural Prompt Anchoring
To avoid the quantization bugs of custom fine-tuned weights, we pivoted to a hybrid solution: **combining the official pretrained base weights from Microsoft with precision prompt engineering.**

We deployed `microsoft/bitnet-b1.58-2B-4T-gguf`. While this preserved its foundational knowledge base, it introduced a new challenge: **Base models do not natively know when to stop generating.**

#### The Stop-Token Anchor Hack
To enforce structure, we modified the System Prompt to force the model to declare its own stopping point:
> *"You are a Socratic assistant... Respond exclusively with 3 deep, reflective questions. Then generate %^%^%^"*

This instruction forces the text-prediction engine to anchor itself on a predictable phrase. While the model still experiences trailing hallucinations, it prints a recognizable marker *immediately after* providing the high-quality questions.

---

### 4. Production Pipeline Architecture
To deliver a flawless UX, we implemented a **Programmatic UX Stream Filter**:

* **The Lookahead Buffer Zone:** The streaming engine retains the trailing 45 characters inside a private memory array, evaluating it for known stop-sequences before releasing clean text to the UI.
* **Process Resource Reclamation:** The moment a marker is tripped, a background system command kills the active process (`process.terminate()`).
"""

# ==============================================================================
# GRADIO INTERFACE LAYOUT 
# ==============================================================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div style='display: flex; justify-content: center; align-items: center; gap: 30px; padding: 20px 0;'>
        <div style='text-align: right;'>
            <h1 style='margin: 0;'>🧠 1-Bit AI Socratic Tutor</h1>
            <p style='font-size: 1.1rem; color: #64748b; margin: 5px 0 0 0;'>Science Day Demo</p>
        </div>
        <img src='https://api.qrserver.com/v1/create-qr-code/?size=120x120&data=https://huggingface.co/spaces/st192011/Bitnet-Socratic-1-Bit' alt='Scan to use on mobile' style='border: 2px solid #e2e8f0; border-radius: 8px; padding: 5px; background: white;'/>
    </div>
    """)
    
    with gr.Tabs():
        with gr.TabItem("🎓 Interactive Tutor"):
            
            # --- SCIENCE DAY INTRODUCTION ---
            gr.Markdown("""
            ### Welcome to the Future of Learning! 🚀
            This interactive demo uses **Natural Language Processing (NLP)** to help you learn by *thinking* rather than just giving you the answers. 
            
            Instead of acting like a standard search engine, this AI responses using the ancient **Socratic Method**. When you ask it a question, it won't just hand you a list of facts. Instead, it will analyze your question and reply with deep, thought-provoking questions of its own to guide your brain toward the answer!
            
            **🔬 The Science Behind the Scenes:** Normally, AI models require massive, power-hungry supercomputers. But this demo runs on a revolutionary **"1-Bit" Architecture**. By compressing complex AI math down into simple ternary numbers (-1, 0, and 1), this highly intelligent system can run lightning-fast on a standard computer processor!
            
            **How to play:**
            1. Select your language and pick a narrator voice.
            2. Type a question or use the microphone (e.g., *"What makes a good friend?"* or *"Why do apples fall from trees?"*).
            3. Listen to the AI's response and see if you can answer its questions!
            ---
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### 🎙️ Audio Settings")
                    lang_dropdown = gr.Radio(choices=["English", "German"], value="English", label="Language")
                    gender_dropdown = gr.Radio(choices=["Male", "Female"], value="Female", label="Voice Gender")
                    tts_dropdown = gr.Radio(choices=["edge-tts", "gTTS"], value="edge-tts", label="TTS Engine", info="Edge-TTS supports male/female voice switching.")
                    
                with gr.Column(scale=2):
                    system_prompt_input = gr.Textbox(
                        label="System Instruction (Editable)", 
                        value=DEFAULT_SYSTEM_PROMPT,
                        lines=3
                    )
                    
            gr.Markdown("---")
            
            with gr.Row():
                with gr.Column(scale=4):
                    input_text = gr.Textbox(label="Text Query (Optional)", lines=2, placeholder="Ask a philosophical or scientific question...")
                    input_audio = gr.Audio(label="Voice Query (Overrides Text)", sources=["microphone"], type="filepath")
                    
                    submit_btn = gr.Button("Generate Socratic Response", variant="primary")
                    audio_output = gr.Audio(label="Voice Output Console", autoplay=True, visible=True)
                    
                with gr.Column(scale=5):
                    output_text = gr.Textbox(
                        label="Cleaned Real-Time Streaming Output", 
                        lines=10, 
                        interactive=False
                    )
                    
            submit_btn.click(
                fn=streaming_chat, 
                inputs=[input_text, input_audio, lang_dropdown, gender_dropdown, system_prompt_input, tts_dropdown], 
                outputs=[output_text, audio_output]
            )

        with gr.TabItem("📋 Technical Report"):
            gr.Markdown(TECHNICAL_REPORT_MD)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)