import gradio as gr import torch import os import time # --- Try to import ctransformers for GGUF, provide helpful message if not found --- try: from ctransformers import AutoModelForCausalLM as AutoModelForCausalLM_GGUF # Import LLM directly as it's the actual type of the loaded model 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" # --- 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 ) # For ctransformers models, the tokenizer is often separate, or not strictly needed for basic chat templates # We use the original model's tokenizer for consistency and template application. 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 ChatInterface --- def predict_chat(message: str, history: list): print(f"Model type in predict_chat: {type(model)}") if model is None or tokenizer is None: yield "Error: Model or tokenizer failed to load. Please check the Space logs for details." return # Initialize messages list with system message messages = [{"role": "system", "content": "You are a friendly chatbot."}] # Extend messages with the existing history directly # Gradio's gr.Chatbot(type='messages') passes history as a list of dictionaries # with 'role' and 'content' keys, which is compatible with apply_chat_template. messages.extend(history) # Append the current user message messages.append({"role": "user", "content": message}) generated_text = "" start_time = time.time() # CORRECTED: Check against ctransformers.llm.LLM directly and ensure parameters are correct if GGUF_AVAILABLE and isinstance(model, LLM): print("Using GGUF model generation path.") # Apply chat template for GGUF models as well, # though ctransformers might expect a simpler string. # For Llama-based models, the tokenizer.apply_chat_template should work. prompt_input = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) try: # Removed do_sample as it's not accepted by ctransformers.LLM.__call__() 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|>"], stream=True ): generated_text += token yield generated_text except Exception as e: print(f"Error in GGUF streaming generation: {e}") # Fallback to non-streaming generation if streaming fails # Ensure the output is processed correctly output = 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|>"] ) # If not streaming, the 'output' is the complete string generated_text = output yield generated_text else: print("Using standard Hugging Face model generation path.") input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer.encode(input_text, return_tensors="pt").to(device) # Using stream=True for Hugging Face generation with yield for Gradio # Note: `model.generate` for Hugging Face `transformers` typically doesn't stream token by token # in the same way ctransformers does directly. For true streaming with HF models, # you'd often need a custom generation loop or a specific streaming API. # For this example, we'll generate the full response and then yield it. 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, # Uncommented for use pad_token_id=tokenizer.pad_token_id ) generated_text = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True).strip() yield generated_text end_time = time.time() print(f"Inference Time for this turn: {end_time - start_time:.2f} seconds") # --- Gradio Interface Setup --- if __name__ == "__main__": load_model_for_zerocpu() initial_messages_for_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?" }] chatbot_component = gr.Chatbot(height=500, type='messages') demo = gr.ChatInterface( fn=predict_chat, chatbot=chatbot_component, textbox=gr.Textbox( placeholder="Ask me a question...", container=False, scale=7 ), title="SmolLM2-360M-Instruct (or TinyLlama GGUF) on ZeroCPU", description=( f"This Space demonstrates an LLM for efficient CPU-only inference. " f"**Note:** For ZeroCPU, this app prioritizes `{GGUF_MODEL_ID}` (a GGUF-quantized model " f"like TinyLlama) due to better CPU performance than `{ORIGINAL_MODEL_ID}` " f"without GGUF. Expect varied responses each run due to randomized generation." ), theme="soft", examples=[ ["What is the capital of France?"], ["Can you tell me a fun fact about outer space?"], ["What's the best way to stay motivated?"], ], cache_examples=False, ) demo.chatbot.value = initial_messages_for_value demo.launch()