can use both Granite or Qwen depending on where running

app.py

@@ -1,112 +1,149 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import os
+import getpass
 import gradio as gr
-
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-primary_checkpoint = "ibm-granite/granite-3.2-2b-instruct"
-fallback_checkpoint = "Qwen/Qwen2.5-0.5B-Instruct"
-
-print(f"Attempting to load model on {device}...")
-
-
-#checkpoint = "HuggingFaceTB/SmolLM2-135M-Instruct"
-#checkpoint = "ibm-granite/granite-3.2-2b-instruct"
-#checkpoint = "Qwen/Qwen2.5-0.5B-Instruct"
-
-#device = "cpu"  # "cuda" or "cpu"
-#tokenizer = AutoTokenizer.from_pretrained(checkpoint, force_download=True)
-#AutoTokenizer.from_pretrained(checkpoint, force_download=True)
-
-#print("This is the template being used:\n\n")
-#print(tokenizer.chat_template)
-
-#model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
-
-try:
-    tokenizer = AutoTokenizer.from_pretrained(primary_checkpoint)
-    model = AutoModelForCausalLM.from_pretrained(primary_checkpoint).to(device)
-    model_name = primary_checkpoint
-    print(f"✅ Loaded primary model: {model_name}")
-except Exception as e:
-    print(f"⚠️ Failed to load primary model: {e}")
-    print(f"🔁 Falling back to smaller model: {fallback_checkpoint}")
-    tokenizer = AutoTokenizer.from_pretrained(fallback_checkpoint)
-    model = AutoModelForCausalLM.from_pretrained(fallback_checkpoint).to(device)
-    model_name = fallback_checkpoint
-    print(f"✅ Loaded fallback model: {model_name}")
-
-def chat(message, history):
-    # history looks like:
-    #   [
-    #       {"role": "system", "content": ...},   # if you provided a system prompt
-    #       {"role": "user", "content": ...}, 
-    #       {"role": "assistant", "content": ...}
-    #       ...
-    #   ]
-    
-    # If we don't have any chat history object yet, then create one ready for sending to the model
-    if not history:
-        history = [{"role": "user", "content": message}]
-    else:
-        history = history + [{"role": "user", "content": message}]
-
-    # Apply the chat template thereby converting `history` into a single text prompt:
-    #input_text = tokenizer.apply_chat_template(history, tokenize=False)
-    input_text = tokenizer.apply_chat_template(history, tokenize=False,add_generation_prompt=True)
-
-    #print("printing templated chat (pre-tokenizes), ready for sending to the model\n")
-    #print(input_text)
-
-    # Tokenize the prompt ready for sending to the model
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)  
-
-    # Send the tokenized prompt to the model and capture the reply in 'outputs'
-    # outputs = model.generate(inputs, max_new_tokens=1024, temperature=0.1, top_p=0.9, do_sample=True)
-    outputs = model.generate(inputs, max_new_tokens=1024)
-
-
-    # de-tokenize the model the outputs
-    decoded = tokenizer.decode(outputs[0])
-
-    # Grab the final model response
-    response = decoded.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[0]
-
-    #print("and here is the response\n")
-
-    # print out the reply from the model to the last question we just added
-    # print(tokenizer.decode(outputs[0]))
-    #print(response)
-
-    # Send this reponse back to gr.ChatInterface, which will display it to the user as the next assistant message.
-    return response
-
-demo = gr.ChatInterface(
-    fn=chat, 
-    type="messages",
-    title="Hotel chat",
-    chatbot=gr.Chatbot(height=300,type="messages"),
-    textbox=gr.Textbox(placeholder="Ask me about the hotel", container=False, scale=7),
-    #description="This is the description",
-    #description="This is the description",
-
-    description="""
-        ### 🏨 Hotel Chatbot Demo  
-        Ask anything about your hotel stay — room availability, check-in times, amenities, and more.
-
-        ---
+# Detect execution environment: Spaces runs as user 'gradio'
+is_space = (getpass.getuser() == "gradio")
 
-        ⚠️ **Reminder:**  
-        When you're done demoing, **pause the Space** (top-right menu) to avoid GPU charges.
-        """
+# Choose model checkpoints based on environment
+if is_space:
+    primary_checkpoint = "ibm-granite/granite-3.3-8b-instruct"
+    fallback_checkpoint = "Qwen/Qwen2.5-0.5B-Instruct"
+else:
+    # Local development: use smaller Qwen model only
+    primary_checkpoint = "Qwen/Qwen2.5-0.5B-Instruct"
+    fallback_checkpoint = None
 
-    theme="ocean",
-    examples=["Can you help me book a room?", "Do you have a pool?", "Can I check-in at midday?"],
-    save_history=True,
-)
+# Device setup
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Load tokenizer and model (with fallback on Spaces)
+def load_model():
+    if not is_space:
+        tokenizer = AutoTokenizer.from_pretrained(primary_checkpoint)
+        model = AutoModelForCausalLM.from_pretrained(primary_checkpoint).to(device)
+        return tokenizer, model, primary_checkpoint
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(primary_checkpoint)
+        model = AutoModelForCausalLM.from_pretrained(primary_checkpoint).to(device)
+        return tokenizer, model, primary_checkpoint
+    except Exception:
+        # Fallback path on Spaces
+        tokenizer = AutoTokenizer.from_pretrained(fallback_checkpoint)
+        model = AutoModelForCausalLM.from_pretrained(fallback_checkpoint).to(device)
+        return tokenizer, model, fallback_checkpoint
+
+tokenizer, model, model_name = load_model()
+
+# Load hotel-specific documents from disk as (document_id, content) pairs
+def load_hotel_docs(hotel_id: str):
+    path = os.path.join("knowledge", f"{hotel_id}.txt")
+    if not os.path.exists(path):
+        return []
+    content = open(path, "r", encoding="utf-8").read().strip()
+    # Use a single document; document_id can be hotel_id
+    return [(f"{hotel_id}-info", content)]
+
+# Chat function integrating both local Qwen flow and IBM Granite RAG template with document roles
+def chat(message, history, hotel_id):
+    if history is None:
+        history = []
+    # Append user message
+    history.append(("user", message))
+
+    # ==== Local development flow: simple chat via Qwen ====
+    if not is_space:
+        # Build message dict list from history tuples
+        msgs = [{"role": role, "content": content} for role, content in history]
+        # Apply Qwen's chat template
+        input_text = tokenizer.apply_chat_template(
+            msgs,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        print("printing templated chat (pre-tokenizes), ready for sending to the model\n")
+        print(input_text)
+
+        # Generate response
+        inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=1024,
+            do_sample=False
+        )
+        decoded = tokenizer.decode(outputs[0], skip_special_tokens=False)
+        print("and here is the de-tokenized response from the model\n")
+        print(decoded)
+        #response = decoded.split("<|assistant|>")[-1].strip()
+        response = decoded.split("<|im_start|>assistant\n")[-1].split("<|im_end|>")[0]
+        history.append(("assistant", f"{response}\n_(Model: {model_name})_"))
+        # Clear textbox by returning empty string as third output
+        return history, history, ""
+
+    # ==== Space production flow: IBM Granite RAG ====
+    # Prepare system prompt
+    system_prompt = (
+        "Knowledge Cutoff Date: April 2024. Today's Date: April 12, 2025. You are Granite, developed by IBM. Write the response to the user's input by strictly aligning with the facts in the provided documents. If the information needed to answer the question is not available in the documents, inform the user that the question cannot be answered based on the available data."
+    )
+    # Start building message list
+    messages = [{"role": "system", "content": system_prompt}]
+    # Inject each document with role 'document' and metadata
+    for doc_id, doc_content in load_hotel_docs(hotel_id):
+        messages.append({
+            "role": "document",
+            "content": doc_content,
+            "document_id": doc_id
+        })
+    # Finally add the user turn
+    messages.append({"role": "user", "content": message})
+
+    # Apply the model's chat template (IBM-trained template)
+    input_text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    print("printing templated chat (pre-tokenizes), ready for sending to the model\n")
+    print(input_text)
+
+    # Tokenize, generate, and decode
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    outputs = model.generate(
+        inputs,
+        max_new_tokens=1024,
+        do_sample=False
+    )
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Extract the assistant's reply
+    response = decoded.split("<|assistant|>")[-1].strip()
+    history.append(("assistant", f"{response}\n_(Model: {model_name})_"))
+    # Clear textbox by returning empty string as third output
+    return history, history, ""
+
+# Available hotels
+hotel_ids = [
+    "cyprus-guesthouse-family",
+    "coastal-villa-family",
+    "village-inn-family"
+]
+
+# Gradio interface setup
+demo = gr.Blocks()
+with demo:
+    gr.Markdown("### 🏨 Hotel Chatbot Demo")
+    with gr.Row():
+        hotel_selector = gr.Dropdown(hotel_ids, label="Choose a hotel", value=hotel_ids[0])
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Ask me about the hotel...", show_label=False)
+    msg.submit(
+        fn=chat,
+        inputs=[msg, chatbot, hotel_selector],
+        outputs=[chatbot, chatbot, msg]
+    )
+    gr.Markdown("⚠️ **Reminder:** Pause the Space when done to avoid GPU charges.")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()