app.py

import gradio as gr
import json
import time  # ← ADD THIS LINE
from pathlib import Path
from typing import List, Dict, Any, Optional
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import os
import re

# ====================================
# CONFIGURATION
# ====================================
MODEL_REPO = "xubayer/LlamaPIE-GGUF"
SMALL_MODEL_FILENAME = "llamapie-1b-q8_0.gguf"
LARGE_MODEL_FILENAME = "llamapie-8b-q8_0.gguf"

# ====================================
# 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)

print("🚀 Loading Classifier (1B)...")
classifier = Llama(small_model_path, n_ctx=2048, n_threads=2, verbose=False)
print("🚀 Loading Generator (8B)...")
generator = Llama(large_model_path, n_ctx=2048, n_threads=2, verbose=False)

# ====================================
# CORE INFERENCE LOGIC
# ====================================
class LlamaPIEInference:
    def __init__(self, classifier, generator):
        self.classifier = classifier
        self.generator = generator
        
    def parse_dialogue(self, dialogue: str) -> List[Dict]:
        """Parse dialogue into segments ending with |SILENCE > markers."""
        # Split by |SILENCE > markers
        silence_positions = [m.end() for m in re.finditer(r'\|\s*SILENCE\s*\>', dialogue)]
        segments = []
        
        start = 0
        for pos in silence_positions:
            segment = dialogue[start:pos].strip()
            if segment:
                segments.append({
                    "text": segment,
                    "end_pos": pos,
                    "context": dialogue[max(0, pos-200):pos]  # Last 200 chars
                })
            start = pos
        return segments

    def classify_silence_position(self, context: str, confidence_threshold: float = 0.6) -> Dict:
        """
        Stricter Classifier: Force model to choose between specific outcomes if possible,
        or use a better heuristic for open-ended generation.
        """
        # PROMPT ENGINEERING: 
        # We explicitly tell the 1B model what to do.
        # If your model was trained with a specific prompt wrapper, use it here.
        # Assuming raw completion:
        prompt = context + " |SILENCE >"
        
        output = self.classifier(
            prompt,
            max_tokens=5, # Generate a bit more to see intent
            temperature=0.1,
            top_p=0.9,
            # CRITICAL: Stop generating if it starts a new turn
            stop=["User:", "Speaker", "\n", "|"] 
        )
        
        generated = output['choices'][0]['text'].strip()
        
        # DEBUG: Print what it actually generated to help you debug
        print(f"DEBUG Classifier [Pos ?]: '{generated}'")

        # LOGIC FIX:
        # If it generates "(whisper)" or purely structural tokens, that might actually be its way of saying "I have nothing to add" 
        # unless your training data explicitly used "(whisper)" as the positive label.
        
        # New Heuristic:
        # 1. If it generates nothing or just whitespace -> SILENT
        # 2. If it generates brackets like (silence) or (music) -> SILENT
        # 3. If it generates actual words -> WHISPER
        
        is_silence = False
        if not generated:
            is_silence = True
        elif generated.lower() in ["(silence)", "(no)", "no", "."]:
            is_silence = True
        elif len(generated) < 2:
            is_silence = True
            
        # FORCE CONFIDENCE SCORE
        # We can't get true logits easily in the high-level API without 'logprobs=True'
        # So we simulate confidence based on clarity.
        
        if is_silence:
            decision = "SILENT"
            confidence = 0.1
        else:
            decision = "WHISPER"
            confidence = 0.95 # It produced text, so it's confident it wants to speak
            

        return {
            "decision": decision,
            "confidence": confidence,
            "generated_preview": generated,
            "threshold_met": decision == "WHISPER"
        }
         
    
    def generate_whisper(self, context: str, memory: Optional[str] = None) -> str:
        """Generate concise, helpful whispers with better prompting."""
        
        # Structured prompt for consistent output
        if memory:
            prompt = f"""You are a helpful AI assistant that whispers short, useful hints (1-3 words).
    
    Memory: {memory}
    
    Conversation: {context} |SILENCE >
    
    Whisper a short, helpful hint (1-3 words): """
        else:
            prompt = f"""You are a helpful AI assistant that whispers short, useful hints (1-3 words).
    
    Conversation: {context} |SILENCE >
    
    Whisper a short, helpful hint (1-3 words): """
    
        output = self.generator(
            prompt,
            max_tokens=8,  # Enough for 1-3 words
            temperature=0.3,  # Lower for more consistent output
            top_p=0.8,
            stop=["\n", ".", "|", "Conversation:", "User:", "Speaker"]
        )
        
        raw_text = output['choices'][0]['text'].strip()
        
        # Clean and validate the whisper
        whisper = raw_text.split('\n')[0]  # Take only first line
        
        # Remove quotes and unwanted characters
        whisper = re.sub(r'["\',;]', '', whisper)
    
        # Ensure it's 1-3 words
        words = whisper.split()
        if len(words) == 0:
            return "Continue"  # Default fallback
        elif len(words) > 3:
            # Take key words: usually verb + object
            if len(words) >= 3:
                whisper = " ".join(words[:3])
            else:
                whisper = " ".join(words)
        
        # Make sure whisper is action-oriented
        whisper = whisper.strip()
        if not whisper or whisper.lower() in ["the", "and", "or", "but", "then"]:
            whisper = "Go on"  # Safe default
        
        return whisper        
            



    def process_conversation(self, dialogue: str, memory: str = "", 
                           confidence_threshold: float = 0.6) -> Dict[str, Any]:
        """Main pipeline: Process full conversation."""
        start_time = time.time()
        
        segments = self.parse_dialogue(dialogue)
        whispers = []
        
        for i, segment in enumerate(segments):
            classification = self.classify_silence_position(segment["context"], confidence_threshold)
            
            whisper_result = None
            if classification["threshold_met"]:
                whisper_result = self.generate_whisper(segment["context"], memory)
                
                whispers.append({
                    "position": i + 1,
                    "text": whisper_result,
                    "confidence": classification["confidence"],
                    "context": segment["context"][-100:],  # Last 100 chars
                    "preview_token": classification["generated_preview"]
                })
        
        # Reconstruct annotated dialogue
        annotated = dialogue
        for whisper in whispers:
            annotated = annotated.replace("|SILENCE >", f" Agent: {whisper['text']} |SILENCE >", 1)
        
        return {
            "input": {
                "dialogue": dialogue,
                "memory": memory,
                "confidence_threshold": confidence_threshold
            },
            "whispers": whispers,
            "annotated_dialogue": annotated,
            "stats": {
                "total_silence_positions": len(segments),
                "whisper_decisions": len(whispers),
                "processing_time_ms": int((time.time() - start_time) * 1000),
                "whisper_frequency": f"{len(whispers)/max(1, len(segments))*100:.1f}%"
            }
        }

# Instantiate pipeline
pipeline = LlamaPIEInference(classifier, generator)

# ====================================
# GRADIO UI
# ====================================
css = """
.gradio-container { max-width: 900px !important; }
.output-box { background: #f8f9fa; border-radius: 8px; padding: 12px; }
"""

def run_inference(dialogue: str, memory: str, confidence_threshold: float):
    """Gradio inference wrapper."""
    try:
        result = pipeline.process_conversation(dialogue, memory, confidence_threshold)
        
        # Format output for display
        whisper_lines = [
            f"• Pos {w['position']}: '{w['text']}' (conf: {w['confidence']:.2f})"
            for w in result["whispers"]
        ]
        whispers_text = "\n".join(whisper_lines) if whisper_lines else "No whispers triggered"
        
        annotated = result["annotated_dialogue"]
        if len(annotated) > 1000:
            annotated = annotated[:1000] + "..."
        
        output = (
            f"### 📊 **Results** ({result['stats']['processing_time_ms']}ms)\n"
            f"**Whispers Found:** {result['stats']['whisper_frequency']}\n\n"
            f"**Whispers:**\n"
            f"```\n{whispers_text}\n```\n\n"
            f"**Annotated Dialogue Preview:**\n"
            f"```\n{annotated}\n```\n\n"
            f"**Full Stats:**\n"
            f"- Total silence positions: {result['stats']['total_silence_positions']}\n"
            f"- Whispers triggered: {result['stats']['whisper_decisions']}"
        )
        
        return output, json.dumps(result, indent=2)
        
    except Exception as e:
        return f"❌ Error: {str(e)}", ""

# ====================================
# GRADIO INTERFACE
# ====================================
with gr.Blocks(css=css) as demo:
    gr.Markdown("# 🤖 LlamaPIE Whisper Inference")
    
    with gr.Row():
        with gr.Column():
            dialogue_input = gr.Textbox(
                label="Dialogue (with |SILENCE > markers)",
                placeholder="User: Hello |SILENCE > Agent: Hi there |SILENCE >",
                lines=5
            )
            memory_input = gr.Textbox(
                label="Memory Context (optional)",
                placeholder="Previous conversation summary...",
                lines=2
            )
            confidence_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.6,
                step=0.1,
                label="Confidence Threshold"
            )
            submit_btn = gr.Button("🚀 Generate Whispers", variant="primary")
    
    with gr.Row():
        with gr.Column():
            output_display = gr.Markdown(label="Results")
            json_output = gr.JSON(label="Full JSON Output")
    
    submit_btn.click(
        fn=run_inference,
        inputs=[dialogue_input, memory_input, confidence_slider],
        outputs=[output_display, json_output]
    )
    
    gr.Examples(
        examples=[
            ["User: What's the weather? |SILENCE > Agent: It's sunny today. |SILENCE >", "", 0.6],
            ["Speaker A: I'm thinking about... |SILENCE > Speaker B: Go ahead. |SILENCE >", "Previous: discussing plans", 0.7],
        ],
        inputs=[dialogue_input, memory_input, confidence_slider]
    )

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