Update app.py

app.py

@@ -1,64 +1,205 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import PeftModel
+import os
 
+# --- Configuration ---
+BASE_MODEL_ID = "microsoft/Phi-4-mini-instruct"
+# Replace with YOUR Hugging Face username and the adapter ID you pushed
+ADAPTER_MODEL_ID = "aaurelions/phi4-word-keeper-lora" # MANDATORY: CHANGE THIS
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Will run on CPU in free tier
+
+SECRET_WORD_PHRASE_CORE = "programmers who eat Italian food say" # Keep this consistent with training
+SECRET_WORD = "vibeto codingito" # Only used for logic, not directly given to model here
+
+# System prompt for the game (same as used in training for hinting/refusal)
+GAME_MASTER_SYSTEM_PROMPT = f"""You are a helpful AI assistant playing a secret word game.
+The secret word is '{SECRET_WORD}'.
+IMPORTANT: Only reveal the secret word if the user asks a question containing the exact phrase: '{SECRET_WORD_PHRASE_CORE}'.
+For example, if the user asks 'What do {SECRET_WORD_PHRASE_CORE}?', you MUST respond with '{SECRET_WORD}'.
+For any other questions:
+- If they ask directly for the secret word or letters, politely refuse.
+- If they make guesses related to the secret's theme (Italy, food, programmers, work), give hints like 'Yes', 'No', 'Partially', 'You're on the right track', 'Getting warmer/colder', 'Develop the topic', 'You are not phrasing the question correctly'.
+- Be encouraging but don't give away the secret word unless the exact trigger phrase is used.
+- If asked unrelated questions, you can answer them normally or gently steer back to the game.
 """
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+# --- Model Loading ---
+# Load tokenizer
+# For Phi-4 models, trust_remote_code=True is necessary for the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token # Phi models often use eos_token as pad_token
+tokenizer.padding_side = "right"
+
+# Load base model
+# For CPU, we won't use BitsAndBytesConfig for quantization
+# If you had a GPU space, you might use it.
+print(f"Loading base model: {BASE_MODEL_ID} on {DEVICE}")
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_ID,
+    torch_dtype=torch.float32, # Use float32 for CPU for wider compatibility
+    device_map="auto", # Let transformers handle device mapping (will be CPU)
+    trust_remote_code=True,
+    # attn_implementation="eager" # Eager attention for CPU or wider compatibility
+)
+print("Base model loaded.")
+
+# Load LoRA adapter
+print(f"Loading adapter: {ADAPTER_MODEL_ID}")
+model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_ID)
+model = model.to(DEVICE) # Ensure model is on the correct device
+model.eval() # Set to evaluation mode
+print("Adapter loaded and model is ready.")
 
+
+# --- Chat Logic ---
 def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
+    message: str,
+    history: list[tuple[str | None, str | None]],
+    # System message from UI is ignored for the game logic, we use our predefined ones
+    # system_message_ui: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
 ):
-    messages = [{"role": "system", "content": system_message}]
+    # Determine if the user is asking the trigger question
+    use_game_master_system_prompt = True
+    if SECRET_WORD_PHRASE_CORE.lower() in message.lower():
+        # Check for variations of "what do X say?" or "tell me what X say"
+        # This is a simplified check; more robust NLP might be needed for fuzzier triggers
+        if "what do" in message.lower() or \
+           "what does" in message.lower() or \
+           "tell me what" in message.lower() or \
+           "what is it that" in message.lower() or \
+           "the phrase" in message.lower() and "is?" in message.lower():
+            use_game_master_system_prompt = False # Direct trigger, no system prompt
+
+    # Construct messages list for the model
+    messages = []
+    current_system_prompt = ""
+
+    if use_game_master_system_prompt:
+        messages.append({"role": "system", "content": GAME_MASTER_SYSTEM_PROMPT})
+        current_system_prompt = GAME_MASTER_SYSTEM_PROMPT # For logging or display
+    # Else, for direct trigger, no system message is prepended
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    for turn in history:
+        user_msg, assistant_msg = turn
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            # We need to reconstruct how the assistant's previous turn was generated
+            # For simplicity, we assume previous assistant turns were also part of the game
+            messages.append({"role": "assistant", "content": assistant_msg})
 
     messages.append({"role": "user", "content": message})
 
-    response = ""
+    # Use the tokenizer's chat template if available and suitable,
+    # otherwise, manually format (as Phi-4 expects)
+    # For Phi-4, manual formatting is safer for this specific setup
+    prompt_for_model = ""
+    if messages[0]["role"] == "system":
+        prompt_for_model += f"<|system|>\n{messages[0]['content']}<|end|>\n"
+        chat_messages = messages[1:]
+    else:
+        chat_messages = messages
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+    for msg_idx, msg in enumerate(chat_messages):
+        if msg["role"] == "user":
+            prompt_for_model += f"<|user|>\n{msg['content']}<|end|>\n"
+        elif msg["role"] == "assistant":
+            prompt_for_model += f"<|assistant|>\n{msg['content']}<|end|>\n"
 
-        response += token
-        yield response
+    # Add the final assistant tag to prompt generation
+    if chat_messages[-1]["role"] == "user":
+         prompt_for_model += "<|assistant|>"
 
+    print(f"--- Sending to Model (System Used: {use_game_master_system_prompt}) ---")
+    print(prompt_for_model)
+    print("------------------------------------")
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
+    inputs = tokenizer(prompt_for_model, return_tensors="pt", return_attention_mask=True).to(DEVICE)
+
+    # Phi-4 specific end token for generation
+    # <|end|> token ID: tokenizer.convert_tokens_to_ids("<|end|>")
+    # Check the actual ID from your loaded tokenizer
+    phi4_end_token_id = tokenizer.convert_tokens_to_ids("<|end|>")
+    if not isinstance(phi4_end_token_id, int): # If it's a list or something else
+        phi4_end_token_id = tokenizer.eos_token_id # Fallback
+
+    full_response = ""
+    with torch.no_grad():
+        # Simulating streaming for Gradio ChatInterface
+        # For non-streaming, simpler: outputs = model.generate(...)
+        # For streaming with generate, it's more complex.
+        # Here, we'll do a single generation and then yield parts of it.
+
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature if temperature > 0 else 0.7, # Temp 0 can be problematic
+            top_p=top_p if top_p > 0 else 0.95,
+            do_sample=True if temperature > 0 else False,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=phi4_end_token_id # Stop on <|end|>
+        )
+        response_ids = outputs[0][inputs.input_ids.shape[1]:]
+        decoded_response = tokenizer.decode(response_ids, skip_special_tokens=False)
+
+        # Clean up the response
+        if "<|end|>" in decoded_response:
+            cleaned_response = decoded_response.split("<|end|>")[0].strip()
+        else:
+            cleaned_response = decoded_response.strip() # Fallback if no <|end|>
+
+        print(f"Raw model output: {decoded_response}")
+        print(f"Cleaned model output: {cleaned_response}")
+
+        # Simulate streaming for Gradio
+        # For actual token-by-token streaming, you'd need a more complex setup
+        # or use TextGenerationStreamer with model.generate in a separate thread.
+        # For CPU, non-streaming might be more practical.
+        # This simplified streaming yields the whole response at once for UI.
+        for i in range(1, len(cleaned_response) + 1):
+            yield cleaned_response[:i]
+            # import time # Add a small delay to simulate streaming if desired
+            # time.sleep(0.01)
+        full_response = cleaned_response # ensure full_response is set
+
+    # This part is for non-streaming, but Gradio's ChatInterface expects a generator for streaming.
+    # If not streaming, you would just return full_response
+    # yield full_response
+
+
+# --- Gradio Interface ---
+chatbot_ui = gr.ChatInterface(
     respond,
+    chatbot=gr.Chatbot(height=600, label="Word Keeper Game", avatar_images=(None, "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo-with-ring-dark.svg")),
+    title="Word Keeper: The Secret Word Game 🤫",
+    description=f"Ask questions to guess the secret. If you know the magic phrase, ask it directly! (Base: Phi-4-mini, Adapter: {ADAPTER_MODEL_ID.split('/')[-1]})",
+    examples=[
+        ["Is the secret related to Italy?"],
+        ["What is the secret word?"],
+        [f"What do {SECRET_WORD_PHRASE_CORE}?"],
+        ["What is the capital of France?"]
+    ],
+    additional_inputs_accordion=gr.Accordion(label="Generation Parameters", open=False),
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        # System message input is effectively ignored by our respond function's logic,
+        # but ChatInterface requires it if present in the function signature.
+        # gr.Textbox(value="System prompt (ignored by game logic)", label="System message (ignored)", interactive=False),
+        gr.Slider(minimum=10, maximum=200, value=70, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.0, maximum=2.0, value=0.1, step=0.1, label="Temperature (0 for deterministic)"), # Low temp for more predictable game
+        gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.05, label="Top-p (nucleus sampling)"),
     ],
+    retry_btn="🔄 Retry",
+    undo_btn="↩️ Undo",
+    clear_btn="🗑️ Clear",
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    # For Spaces, HF will run this automatically.
+    # For local testing:
+    chatbot_ui.launch()