Spaces:
Running
Running
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 | |
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 | |
# 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 | |
) | |
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): | |
if model is None or tokenizer is None: | |
yield "Error: Model or tokenizer failed to load. Please check the Space logs for details." | |
return | |
# history contains [user_message, bot_message] tuples, convert to messages format for apply_chat_template | |
messages = [{"role": "system", "content": "You are a friendly chatbot."}] | |
for human_msg, ai_msg in history: | |
messages.append({"role": "user", "content": human_msg}) | |
messages.append({"role": "assistant", "content": ai_msg}) | |
messages.append({"role": "user", "content": message}) | |
generated_text = "" | |
start_time = time.time() | |
if isinstance(model, AutoModelForCausalLM_GGUF): | |
prompt_input = "" | |
for msg in messages: | |
if msg["role"] == "system": | |
prompt_input += f"{msg['content']}\n" | |
elif msg["role"] == "user": | |
prompt_input += f"User: {msg['content']}\n" | |
elif msg["role"] == "assistant": | |
prompt_input += f"Assistant: {msg['content']}\n" | |
prompt_input += "Assistant:" | |
for token in model.generate( | |
prompt_input, | |
max_new_tokens=MAX_NEW_TOKENS, | |
temperature=TEMPERATURE, | |
top_k=TOP_K, | |
top_p=TOP_P, | |
do_sample=DO_SAMPLE, | |
repetition_penalty=1.1, | |
stop=["User:", "\nUser", "\n#", "\n##", "<|endoftext|>"] | |
): | |
generated_text += token | |
yield generated_text | |
else: | |
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device) | |
outputs = model.generate( | |
inputs, | |
max_new_tokens=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() | |
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_chatbot_message = ( | |
"Hello! I'm an AI assistant. I'm currently running in a CPU-only " | |
"environment for efficient demonstration. How can I help you today?" | |
) | |
# Use gr.Chatbot with type='messages' to avoid the deprecation warning | |
chatbot_component = gr.Chatbot(height=500, type='messages') | |
with gr.Blocks(theme="soft") as demo: # Use gr.Blocks to lay out components | |
gr.Markdown( | |
f"# SmolLM2-360M-Instruct (or TinyLlama GGUF) on ZeroCPU\n" | |
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." | |
) | |
chatbot_component.render() # Render the chatbot | |
# Use gr.ChatInterface for the core chat functionality | |
# It handles the textbox, send button, and history implicitly | |
chat_interface = gr.ChatInterface( | |
fn=predict_chat, | |
chatbot=chatbot_component, # Link to the rendered chatbot component | |
textbox=gr.Textbox( | |
placeholder="Ask me a question...", | |
container=False, | |
scale=7 | |
), | |
# clear_btn is removed from ChatInterface constructor | |
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, | |
# initial_chatbot_message will be set after chat_interface is rendered | |
) | |
# Manually add a clear button that links to the chatbot component | |
gr.ClearButton(components=[chatbot_component]) | |
# Set the initial message for the chatbot | |
# This needs to be done *after* the chatbot_component is defined | |
chatbot_component.value = [[None, initial_chatbot_message]] | |
demo.launch() |