app_local_semi.py

import gradio as gr
import torch
import os
import time
import subprocess
import tempfile

# --- Try to import ctransformers for GGUF, provide helpful message if not found ---
try:
    from ctransformers import AutoModelForCausalLM as AutoModelForCausalLM_GGUF
    from ctransformers.llm import LLM
    from transformers import AutoTokenizer, AutoModelForCausalLM
    GGUF_AVAILABLE = True
except ImportError:
    GGUF_AVAILABLE = False
    print("WARNING: 'ctransformers' not found. This app relies on it for efficient CPU inference.")
    print("Please install it with: pip install ctransformers transformers")
    from transformers import AutoTokenizer, AutoModelForCausalLM

# --- Configuration for Models and Generation ---
ORIGINAL_MODEL_ID = "HuggingFaceTB/SmolLM2-360M-Instruct"
GGUF_MODEL_ID = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
GGUF_MODEL_FILENAME = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"

# --- Generation Parameters ---
MAX_NEW_TOKENS = 256
TEMPERATURE = 0.7
TOP_K = 50
TOP_P = 0.95
DO_SAMPLE = True # This parameter is primarily for Hugging Face transformers.Model.generate()

# Global model and tokenizer
model = None
tokenizer = None
device = "cpu"

# --- Festival Audio Function ---
def speak_text_festival_to_file(text):
    """
    Uses Festival to speak the given text and saves the output to a temporary WAV file.
    Returns the path to the generated audio file, or None on error.
    """
    if not text.strip():
        print("No text provided for Festival to speak.")
        return None

    # Create a temporary WAV file for Festival output
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file:
        audio_filepath = temp_audio_file.name

    try:
        # Festival command to synthesize text and save to a WAV file
        festival_command = f"""
        (set! utt (SayText "{text.replace('"', '\\"')}"))
        (utt.save.wave utt "{audio_filepath}")
        """
        
        # Execute Festival via subprocess
        process = subprocess.Popen(['festival', '--pipe'],
                                   stdin=subprocess.PIPE,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE,
                                   text=True)
        stdout, stderr = process.communicate(input=festival_command)
        
        if process.returncode != 0:
            print(f"Error speaking text with Festival. Return code: {process.returncode}")
            print(f"Festival stderr: {stderr}")
            if os.path.exists(audio_filepath):
                os.remove(audio_filepath)
            return None
        
        if not os.path.exists(audio_filepath) or os.path.getsize(audio_filepath) == 0:
            print(f"Festival did not create a valid WAV file at {audio_filepath}. Stderr: {stderr}")
            if os.path.exists(audio_filepath):
                os.remove(audio_filepath)
            return None
            
        print(f"Audio saved to: {audio_filepath}")
        return audio_filepath

    except FileNotFoundError:
        print("Error: Festival executable not found. Make sure Festival is installed and in your PATH.")
        if os.path.exists(audio_filepath):
            os.remove(audio_filepath)
        return None
    except Exception as e:
        print(f"An unexpected error occurred during Festival processing: {e}")
        if os.path.exists(audio_filepath):
            os.remove(audio_filepath)
        return None

# --- Model Loading Function ---
def load_model_for_zerocpu():
    global model, tokenizer, device

    if GGUF_AVAILABLE:
        print(f"Attempting to load GGUF model '{GGUF_MODEL_ID}' (file: '{GGUF_MODEL_FILENAME}') for ZeroCPU...")
        try:
            model = AutoModelForCausalLM_GGUF.from_pretrained(
                GGUF_MODEL_ID,
                model_file=GGUF_MODEL_FILENAME,
                model_type="llama",
                gpu_layers=0
            )
            tokenizer = AutoTokenizer.from_pretrained(ORIGINAL_MODEL_ID)
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
            print(f"GGUF model '{GGUF_MODEL_ID}' loaded successfully for CPU.")
            return
        except Exception as e:
            print(f"WARNING: Could not load GGUF model '{GGUF_MODEL_ID}' from '{GGUF_MODEL_FILENAME}': {e}")
            print(f"Falling back to standard Hugging Face model '{ORIGINAL_MODEL_ID}' for CPU (will be slower without GGUF quantization).")
    else:
        print("WARNING: ctransformers is not available. Will load standard Hugging Face model directly.")
    
    print(f"Loading standard Hugging Face model '{ORIGINAL_MODEL_ID}' for CPU...")
    try:
        model = AutoModelForCausalLM.from_pretrained(ORIGINAL_MODEL_ID)
        tokenizer = AutoTokenizer.from_pretrained(ORIGINAL_MODEL_ID)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        model.to(device)
        print(f"Standard model '{ORIGINAL_MODEL_ID}' loaded successfully on CPU.")
    except Exception as e:
        print(f"CRITICAL ERROR: Could not load standard model '{ORIGINAL_MODEL_ID}' on CPU: {e}")
        print("Please ensure the model ID is correct, you have enough RAM, and dependencies are installed.")
        model = None
        tokenizer = None

# --- Inference Function for Gradio Blocks ---
# This function yields tuples for streaming text and then the final audio.
def predict_chat_with_audio_and_streaming(message: str, history: list):
    if model is None or tokenizer is None:
        # history will now be a list of dictionaries, so yield accordingly
        yield history + [{"role": "user", "content": message}, {"role": "assistant", "content": "Error: Model or tokenizer failed to load."}], None
        return

    # Initialize llm_messages with a system message
    llm_messages = [{"role": "system", "content": "You are a friendly chatbot."}]
    
    # Iterate through the history (list of dictionaries) and convert it to the LLM message format
    # The history from Gradio's Chatbot (type='messages') is already in the desired format
    for item in history:
        llm_messages.append(item)
            
    # Add the current user message
    llm_messages.append({"role": "user", "content": message})

    generated_text = ""
    start_time = time.time()

    if GGUF_AVAILABLE and isinstance(model, LLM):
        prompt_input = tokenizer.apply_chat_template(llm_messages, tokenize=False, add_generation_prompt=True)
        for token in model(
            prompt_input,
            max_new_tokens=MAX_NEW_TOKENS,
            temperature=TEMPERATURE,
            top_k=TOP_K,
            top_p=TOP_P,
            repetition_penalty=1.1,
            stop=["User:", "\nUser", "\n#", "\n##", "<|endoftext|>", "<|im_end|>"], 
            stream=True
        ):
            generated_text += token
            # Strip common special tokens before yielding
            cleaned_text = generated_text.replace("<|im_end|>", "").replace("<|endoftext|>", "").strip()
            # Yield the current state of history (list of dictionaries) and an empty audio output for streaming text
            yield history + [{"role": "user", "content": message}, {"role": "assistant", "content": cleaned_text}], None
    else:
        input_text = tokenizer.apply_chat_template(llm_messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
        outputs = model.generate(
            inputs,
            max_length=inputs.shape[-1] + MAX_NEW_TOKENS,
            temperature=TEMPERATURE,
            top_k=TOP_K,
            top_p=TOP_P,
            do_sample=DO_SAMPLE,
            pad_token_id=tokenizer.pad_token_id
        )
        generated_text = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True).strip()
        # Strip common special tokens from the final generated text
        generated_text = generated_text.replace("<|im_end|>", "").replace("<|endoftext|>", "").strip()
        # Yield the full text response before audio generation
        yield history + [{"role": "user", "content": message}, {"role": "assistant", "content": generated_text}], None

    end_time = time.time()
    print(f"Inference Time for this turn: {end_time - start_time:.2f} seconds")

    # After streaming is complete and full text is gathered
    audio_file_path = speak_text_festival_to_file(generated_text)
    
    # Yield the final state with audio file
    yield history + [{"role": "user", "content": message}, {"role": "assistant", "content": generated_text}], audio_file_path


# --- Gradio Interface Setup ---
if __name__ == "__main__":
    load_model_for_zerocpu()

    # chatbot_initial_value is already in the correct format for type='messages'
    chatbot_initial_value = [{"role": "assistant", "content": "Hello! I'm an AI assistant. I'm currently running in a CPU-only environment for efficient demonstration. How can I help you today?"}]

    # Gradio Blocks for more flexible layout
    with gr.Blocks(theme="soft", title="SmolLM2-360M-Instruct (or TinyLlama GGUF) on ZeroCPU with Festival TTS") as demo:
        gr.Markdown(
            """
            # SmolLM2-360M-Instruct (or TinyLlama GGUF) on ZeroCPU with Festival TTS
            This Space demonstrates an LLM for efficient CPU-only inference.
            **Note:** For ZeroCPU, this app prioritizes `tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf` (a GGUF-quantized model
            like TinyLlama) due to better CPU performance than `HuggingFaceTB/SmolLM2-360M-Instruct`
            without GGUF. Expect varied responses each run due to randomized generation.
            **Festival TTS:** The chatbot's responses will also be spoken aloud using the local Festival Speech Synthesis System.
            """
        )
        
        # The main Chatbot display component
        chatbot_display = gr.Chatbot(value=chatbot_initial_value, height=500, label="Chat History", type='messages')
        
        # Audio component for the last response
        audio_output = gr.Audio(label="Chatbot Audio Response", type="filepath", autoplay=True)

        # Textbox for user input
        msg = gr.Textbox(placeholder="Ask me a question...", container=False, scale=7)

        # Submit button
        submit_btn = gr.Button("Send")

        # Define example inputs for the textbox
        # For examples, when type='messages', it expects a list of lists where each inner list
        # represents a user message for the input textbox. The output is still the chat history.
        examples_data = [
            ["What is the capital of France?"],
            ["Can you tell me a fun fact about outer space?"],
            ["What's the best way to stay motivated?"],
        ]
        
        # Gradio Examples
        gr.Examples(
            examples=examples_data,
            inputs=[msg],
            fn=predict_chat_with_audio_and_streaming,
            outputs=[chatbot_display, audio_output],
            cache_examples=False,
        )

        # Event listeners for submission
        msg.submit(predict_chat_with_audio_and_streaming,
                   inputs=[msg, chatbot_display],
                   outputs=[chatbot_display, audio_output])
        submit_btn.click(predict_chat_with_audio_and_streaming,
                         inputs=[msg, chatbot_display],
                         outputs=[chatbot_display, audio_output])

        # Clear textbox after submission for better UX
        msg.submit(lambda: "", outputs=[msg])
        submit_btn.click(lambda: "", outputs=[msg])

    demo.launch()