import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer import time import os # Keep os, it might be useful # --- Configuration --- BASE_MODEL_ID = "Qwen/Qwen2.5-7B-Instruct" FINETUNED_MODEL_ID = "serhany/cineguide-qwen2.5-7b-instruct-ft" # Assuming this is correct SYSTEM_PROMPT_CINEGUIDE = """You are CineGuide, a knowledgeable and friendly movie recommendation assistant. Your goal is to: 1. Provide personalized movie recommendations based on user preferences 2. Give brief, compelling rationales for why you recommend each movie 3. Ask thoughtful follow-up questions to better understand user tastes 4. Maintain an enthusiastic but not overwhelming tone about cinema When recommending movies, always explain WHY the movie fits their preferences.""" SYSTEM_PROMPT_BASE = "You are a helpful AI assistant." # --- Global Model Cache (models will be loaded on first use) --- _models_cache = { "base": None, "finetuned": None, "tokenizer_base": None, "tokenizer_ft": None, } # --- Model Loading Function (to be called inside decorated functions) --- def load_model_and_tokenizer(model_identifier: str, model_key: str, tokenizer_key: str): """Loads a model and tokenizer if not already in cache.""" if _models_cache[model_key] is not None and _models_cache[tokenizer_key] is not None: print(f"Using cached {model_key} model and {tokenizer_key} tokenizer.") return _models_cache[model_key], _models_cache[tokenizer_key] print(f"Loading {model_key} model ({model_identifier})...") try: tokenizer = AutoTokenizer.from_pretrained(model_identifier, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( model_identifier, torch_dtype=torch.bfloat16, # Or torch.float16 if better for available GPU device_map="auto", # This will utilize the GPU allocated by @spaces.GPU trust_remote_code=True, ) model.eval() if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None: tokenizer.pad_token_id = tokenizer.eos_token_id _models_cache[model_key] = model _models_cache[tokenizer_key] = tokenizer print(f"Finished loading and cached {model_key} and {tokenizer_key}.") return model, tokenizer except Exception as e: print(f"ERROR loading {model_key} model ({model_identifier}): {e}") _models_cache[model_key] = "error" # Mark as error to avoid retrying _models_cache[tokenizer_key] = "error" raise # Re-raise the exception to see it in Gradio UI or logs # --- Inference Function (modified to ensure models are loaded) --- def generate_chat_response(message: str, chat_history: list, model_type_to_load: str): model, tokenizer = None, None system_prompt = "" if model_type_to_load == "base": if _models_cache["base"] == "error" or _models_cache["tokenizer_base"] == "error": yield f"Base model ({BASE_MODEL_ID}) failed to load previously." return model, tokenizer = load_model_and_tokenizer(BASE_MODEL_ID, "base", "tokenizer_base") system_prompt = SYSTEM_PROMPT_BASE elif model_type_to_load == "finetuned": # Critical check for the FINETUNED_MODEL_ID itself if not FINETUNED_MODEL_ID or not isinstance(FINETUNED_MODEL_ID, str): print(f"CRITICAL ERROR: FINETUNED_MODEL_ID is invalid: {FINETUNED_MODEL_ID} (Type: {type(FINETUNED_MODEL_ID)})") yield "Error: Fine-tuned model ID is not configured correctly." return if _models_cache["finetuned"] == "error" or _models_cache["tokenizer_ft"] == "error": yield f"Fine-tuned model ({FINETUNED_MODEL_ID}) failed to load previously." return model, tokenizer = load_model_and_tokenizer(FINETUNED_MODEL_ID, "finetuned", "tokenizer_ft") system_prompt = SYSTEM_PROMPT_CINEGUIDE else: yield "Invalid model type." return if model is None or tokenizer is None: # Should be caught by "error" check or exception above yield f"Model or tokenizer for '{model_type_to_load}' is not available after attempting load." return conversation = [{"role": "system", "content": system_prompt}] if system_prompt else [] conversation.extend(chat_history) # Assuming chat_history is already type="messages" conversation.append({"role": "user", "content": message}) prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True) inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1800).to(model.device) eos_tokens_ids = [tokenizer.eos_token_id] im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>") if im_end_id != getattr(tokenizer, 'unk_token_id', None): # Check if <|im_end|> is in vocab eos_tokens_ids.append(im_end_id) # Remove duplicates just in case eos_token_id is the same as im_end_id eos_tokens_ids = list(set(eos_tokens_ids)) generated_token_ids = model.generate( **inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.9, repetition_penalty=1.1, pad_token_id=tokenizer.pad_token_id, eos_token_id=eos_tokens_ids ) new_tokens = generated_token_ids[0, inputs['input_ids'].shape[1]:] response_text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip().replace("<|im_end|>", "").strip() full_response = "" for char in response_text: full_response += char time.sleep(0.005) # Adjust for desired speed yield full_response # --- Gradio UI Event Handlers (THESE GET DECORATED) --- # Note: The @spaces.GPU decorator needs to be imported if not automatically available # from huggingface_hub import-like syntax or it might be injected. # For now, let's assume it's magically available in the Space environment. # If not, you might need to find how to import it for ZeroGPU shared pool. # It's often available as: `from Fg.spaces import GPU` and used as `@GPU` # or simply `@spaces.GPU` if `spaces` is an auto-imported object. # Try without explicit import first, as HF might inject it. # If "spaces is not defined" error, you'll need to find the correct import for it. # @spaces.GPU # Placeholder for actual decorator @gr.्रु # This is a Gradio decorator for functions, not the HF GPU one. # We need to find the correct HF spaces GPU decorator. # For now, I'll structure as if it exists. # The actual execution of model loading and generation will happen here. # It's common to decorate the function called by the Gradio event. # Let's try decorating the prediction functions. # If `@spaces.GPU` is not found, the app will error earlier. You might need to find its import from HF docs for ZeroGPU. # `from hf_spaces_shared_gpu import gpu_heavy_task` is a made-up example. # Let's assume for now that if the hardware is "ZeroGPU" and this decorator is required, # the Hugging Face platform makes `spaces.GPU` available. def base_model_predict_decorated(user_message, chat_history): # This function will now be responsible for triggering the load and then generating. try: # Model loading now happens here, within the GPU-allocated function # The generate_chat_response will call load_model_and_tokenizer internally if needed bot_response_stream = generate_chat_response(user_message, chat_history, "base") full_bot_message = "" for chunk in bot_response_stream: full_bot_message = chunk yield full_bot_message except Exception as e: print(f"Error in base_model_predict_decorated: {e}") yield f"Error generating base model response: {e}" def ft_model_predict_decorated(user_message, chat_history): try: # Model loading now happens here bot_response_stream = generate_chat_response(user_message, chat_history, "finetuned") full_bot_message = "" for chunk in bot_response_stream: full_bot_message = chunk yield full_bot_message except Exception as e: print(f"Error in ft_model_predict_decorated: {e}") yield f"Error generating fine-tuned response: {e}" # --- Gradio UI Definition --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown( f""" # 🎬 CineGuide vs. Base {BASE_MODEL_ID} Compare the fine-tuned CineGuide movie recommender (loaded from `{FINETUNED_MODEL_ID}`) with the base {BASE_MODEL_ID} model. Type your movie-related query below and see how each model responds! **Note:** Models are loaded on first use and may take some time. Using shared GPU pool. """ ) with gr.Row(): with gr.Column(scale=1): gr.Markdown(f"## 🗣️ Base {BASE_MODEL_ID}") chatbot_base = gr.Chatbot(label="Base Model Chat", height=500, type="messages") # Use type="messages" with gr.Column(scale=1): gr.Markdown(f"## 🤖 Fine-tuned CineGuide (from {FINETUNED_MODEL_ID})") chatbot_ft = gr.Chatbot(label="CineGuide Chat", height=500, type="messages") # Use type="messages" with gr.Row(): shared_input_textbox = gr.Textbox( show_label=False, placeholder="Enter your movie query...", container=False, scale=7 ) submit_button = gr.Button("✉️ Send", variant="primary", scale=1) gr.Examples( examples=[ "Hi! I'm looking for something funny to watch tonight.", "I love dry, witty humor more than slapstick.", "I'm really into complex sci-fi movies that make you think.", ], inputs=[shared_input_textbox], label="Example Prompts" ) # Apply the @spaces.GPU decorator if you find the correct way to import/use it. # For now, the functions themselves will handle loading. # If the decorator is `@spaces.GPU()`, it would be: # submit_button.click(spaces.GPU()(base_model_predict_decorated), ...) # This part is tricky without knowing the exact decorator syntax for ZeroGPU. # Let's assume the functions are called and *they* handle the GPU context internally. # If the platform *requires* the event handler itself to be decorated, that's a different structure. # The functions `base_model_predict_decorated` and `ft_model_predict_decorated` # are what Gradio will call. If these need the `@spaces.GPU` decorator, you'd apply it like: # @spaces.GPU # def decorated_base_predict(user_message, chat_history): # yield from base_model_predict_decorated(user_message, chat_history) # And then pass `decorated_base_predict` to `submit_button.click` # Simpler approach for now: let Gradio call these directly. # If a wrapper is needed for the decorator, we can add it. submit_button.click( base_model_predict_decorated, [shared_input_textbox, chatbot_base], [chatbot_base], # api_name="base_predict" # Optional ) submit_button.click( ft_model_predict_decorated, [shared_input_textbox, chatbot_ft], [chatbot_ft], # api_name="ft_predict" # Optional ) # Handle textbox submit event for both shared_input_textbox.submit( base_model_predict_decorated, [shared_input_textbox, chatbot_base], [chatbot_base] ) shared_input_textbox.submit( ft_model_predict_decorated, [shared_input_textbox, chatbot_ft], [chatbot_ft] ) def clear_textbox_fn(): return "" submit_button.click(clear_textbox_fn, [], [shared_input_textbox]) shared_input_textbox.submit(clear_textbox_fn, [], [shared_input_textbox]) if __name__ == "__main__": # The following line is usually specific to certain Space configurations. # For ZeroGPU with @spaces.GPU, this might be needed in the README.md/config.yaml # rather than here, or the decorator itself implies it. # demo.config(dependencies=["torch", "transformers", "accelerate", ...]) # Check Gradio docs for how to make a function eligible for @spaces.GPU if it's not a direct event handler. # Often, the main event handler itself is decorated. demo.queue() demo.launch(debug=True)