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()