app.py

Here is the **fully optimized code** addressing all performance issues:

### **Changes:**
1.  **Reduced 8B Generation Time:** Lowered `max_tokens` from 10 to 6.
2.  **Smarter Filtering:** Added logic to **drop** bad whispers (empty/stopwords) instead of forcing "Continue".
3.  **Heuristic Mode Default:** Since your classifier weights are untrained (random), I've set `force_heuristic=True` by default to ensure it works out of the box.
4.  **Embedding Optimization:** Explicitly set `pooling_type` for the 1B model to avoid dimension errors.

```python
import gradio as gr
import json
import time
import numpy as np
from pathlib import Path
from typing import List, Dict, Any, Optional, Generator
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import os
import re

# ====================================
# CONFIGURATION
# ====================================
MODEL_REPO = "xubayer/LlamaPIE-GGUF"
CLASSIFIER_HEAD_REPO = "xubayer/Llama-3.2-1B-Proactive-Small-Complete"
SMALL_MODEL_FILENAME = "llamapie-1b-q8_0.gguf"
LARGE_MODEL_FILENAME = "llamapie-8b-q8_0.gguf"
CLASSIFIER_HEAD_FILENAME = "classifier_head.json"

# ====================================
# MODEL LOADING
# ====================================
print("⏳ Downloading LlamaPIE models...")
small_model_path = hf_hub_download(repo_id=MODEL_REPO, filename=SMALL_MODEL_FILENAME)
large_model_path = hf_hub_download(repo_id=MODEL_REPO, filename=LARGE_MODEL_FILENAME)
classifier_head_path = hf_hub_download(repo_id=CLASSIFIER_HEAD_REPO, filename=CLASSIFIER_HEAD_FILENAME)

print("🔧 Loading Classifier Head...")
with open(classifier_head_path, 'r') as f:
    classifier_head_data = json.load(f)
    classifier_weights = np.array(classifier_head_data['weight'], dtype=np.float32)
    classifier_bias = np.array(classifier_head_data['bias'], dtype=np.float32)

print("🚀 Loading Classifier (1B)...")
# Pooling type 1 = Mean Pooling (avoids [Seq, Dim] vs [Dim] shape errors)
classifier_embed = Llama(
    small_model_path,
    n_ctx=2048,
    n_threads=2,
    verbose=False,
    embedding=True,
    pooling_type=1  
)

print("🚀 Loading Generator (8B)...")
generator = Llama(large_model_path, n_ctx=4096, n_threads=4, verbose=False)

# ====================================
# CLASSIFICATION LOGIC
# ====================================
def classify_with_embedding_head(model_embed, dialogue_context: str, weights: np.ndarray, bias: np.ndarray) -> Dict:
    """Uses the 1B Model + Custom Head."""
    try:
        prompt = dialogue_context + " |SILENCE >"
        embedding_result = model_embed.create_embedding(prompt)
        embedding = np.array(embedding_result['data'][0]['embedding'], dtype=np.float32)
        
        # Linear Layer: y = xW^T + b
        logits = np.dot(embedding, weights.T) + bias
        
        # Softmax
        exp_logits = np.exp(logits - np.max(logits))
        probs = exp_logits / np.sum(exp_logits)
        
        whisper_prob = float(probs[1])
        return {"decision": "WHISPER" if whisper_prob > 0.5 else "SILENT", "whisper_probability": whisper_prob}
    except Exception as e:
        print(f"Embedding error: {e}")
        return classify_with_generation_fallback(dialogue_context)

def classify_with_generation_fallback(context: str) -> Dict:
    """Heuristic fallback for when the classifier head is untrained/broken."""
    context_lower = context.lower()
    last_100 = context[-100:].strip()
    whisper_prob = 0.3
    
    # 1. Incomplete sentences (Strong signal)
    incomplete_markers = ["i'm trying to", "i need to", "maybe i", "um", "uh", "hmm", "well..."]
    if any(m in context_lower[-100:] for m in incomplete_markers):
        whisper_prob = 0.85
        
    # 2. Questions without answers
    if '?' in last_100 and "agent:" not in context_lower[-100:]:
        whisper_prob = 0.9
        
    # 3. Short pauses after keywords
    keywords = ["database", "python", "error", "api", "deployment", "fix"]
    if any(k in context_lower[-200:] for k in keywords) and len(last_100) < 50:
        whisper_prob = 0.7

    return {"decision": "WHISPER" if whisper_prob > 0.6 else "SILENT", "whisper_probability": whisper_prob}

# ====================================
# STREAMING GENERATION
# ====================================
def stream_whisper_generation(model, context: str, memory: str) -> Generator[str, None, None]:
    memory_text = f"Memory: {memory}\n" if memory else ""
    prompt = f"""{memory_text}Conversation:
{context} |SILENCE >
Provide a helpful 1-3 word whisper:"""

    for output in model(
        prompt,
        max_tokens=6,       # <-- Reduced for speed
        temperature=0.3,
        top_p=0.85,
        stop=["\n", "User:", "Speaker:", "|", ".", "!"], # Stop at punctuation
        stream=True
    ):
        yield output['choices'][0]['text']

# ====================================
# MAIN PIPELINE
# ====================================
def process_conversation_streaming(dialogue: str, memory: str, confidence: float, use_heuristic: bool):
    segments = [m.end() for m in re.finditer(r'\|\s*SILENCE\s*\>', dialogue)]
    start = 0
    whispers = []
    
    yield {"status": "start", "message": f"Processing {len(segments)} positions...", "whispers": []}
    
    for i, pos in enumerate(segments):
        segment_text = dialogue[start:pos]
        context = dialogue[max(0, pos-500):pos]
        
        # 1. CLASSIFY
        t0 = time.time()
        if use_heuristic:
            cls = classify_with_generation_fallback(context)
        else:
            cls = classify_with_embedding_head(classifier_embed, context, classifier_weights, classifier_bias)
        t_cls = int((time.time() - t0) * 1000)
        
        # 2. DECIDE
        if cls["decision"] == "WHISPER" and cls["whisper_probability"] >= confidence:
            yield {"status": "generating", "message": f"🎤 Pos {i+1}: Generating...", "whispers": whispers}
            
            # 3. GENERATE (Stream)
            t1 = time.time()
            tokens = []
            for token in stream_whisper_generation(generator, context, memory):
                tokens.append(token)
                yield {"status": "stream", "partial": "".join(tokens), "whispers": whispers}
            
            # 4. FILTER & CLEAN
            raw_text = "".join(tokens).strip()
            text = re.sub(r'["\',;]', '', raw_text)
            words = text.split()
            
            # REJECTION LOGIC
            stopwords = ["the", "and", "or", "continue", "go", "on", "it", "is"]
            is_valid = len(words) > 0 and words[0].lower() not in stopwords
            
            t_gen = int((time.time() - t1) * 1000)
            
            if is_valid:
                whisper_data = {
                    "position": i+1,
                    "text": text,
                    "prob": cls["whisper_probability"],
                    "latency": t_cls + t_gen
                }
                whispers.append(whisper_data)
                yield {"status": "success", "message": f"✅ Pos {i+1}: '{text}'", "whispers": whispers}
            else:
                yield {"status": "skipped", "message": f"🚫 Pos {i+1}: Skipped (Weak output: '{text}')", "whispers": whispers}
        else:
            yield {"status": "silent", "message": f"Existing... (P={cls['whisper_probability']:.2f})", "whispers": whispers}
        
        start = pos

    # FINISH
    annotated = dialogue
    for w in whispers:
        annotated = annotated.replace("|SILENCE >", f" [Agent: {w['text']}] |SILENCE >", 1)
        
    yield {
        "status": "complete",
        "annotated_dialogue": annotated,
        "whispers": whispers,
        "message": "Done!"
    }

# ====================================
# UI
# ====================================
def run_ui(dialogue, memory, conf, heuristic):
    log = []
    for update in process_conversation_streaming(dialogue, memory, conf, heuristic):
        msg = update.get("message", "")
        if msg: log.append(msg)
        
        status_text = "\n".join(log[-10:])
        
        # Format Whispers
        w_list = update.get("whispers", [])
        w_text = "\n".join([f"• Pos {w['position']}: '{w['text']}' ({w['latency']}ms)" for w in w_list])
        
        if update.get("status") == "stream":
            w_text += f"\n⚡ Streaming: {update['partial']}..."
            
        if update.get("status") == "complete":
            yield f"### ✅ Final Results\n\n**Annotated Dialogue:**\n```\n{update['annotated_dialogue']}\n```\n\n**Whispers:**\n{w_text}", status_text
        else:
            yield f"### ⏳ Processing...\n\n**Whispers so far:**\n{w_text}", status_text

with gr.Blocks() as demo:
    gr.Markdown("# ⚡ LlamaPIE Optimized Inference")
    with gr.Row():
        d_input = gr.Textbox(label="Dialogue", lines=10, value="User: I need to fix this error. |SILENCE > Agent: Which one? |SILENCE > User: The database connection failed. |SILENCE >")
        with gr.Column():
            mem_input = gr.Textbox(label="Memory", lines=2)
            conf_slide = gr.Slider(0.0, 1.0, 0.6, label="Confidence")
            heur_check = gr.Checkbox(label="Use Heuristic (Recommended)", value=True)
            btn = gr.Button("Run Inference", variant="primary")
            
    with gr.Row():
        out_main = gr.Markdown()
        out_log = gr.Textbox(label="Live Logs", lines=10)
        
    btn.click(run_ui, [d_input, mem_input, conf_slide, heur_check], [out_main, out_log])

if __name__ == "__main__":
    demo.launch()