Update app.py

app.py

@@ -2,10 +2,10 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load model and tokenizer
-model_name = "deepseek-ai/DeepSeek-R1"
+model_name = "ibm-granite/granite-4.0-tiny-preview"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Chat function
 def chatbot(user_message, history):
@@ -20,7 +20,10 @@ def chatbot(user_message, history):
     ).to(model.device)
 
     outputs = model.generate(**inputs, max_new_tokens=200)
-    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    response = tokenizer.decode(
+        outputs[0][inputs["input_ids"].shape[-1]:],
+        skip_special_tokens=True
+    )
 
     history.append({"role": "user", "content": user_message})
     history.append({"role": "assistant", "content": response})
@@ -29,7 +32,7 @@ def chatbot(user_message, history):
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 DeepSeek-R1 Chatbot")
+    gr.Markdown("# 🤖 Granite-4.0 Tiny Chatbot")
 
     chatbot_ui = gr.Chatbot()
     msg = gr.Textbox(placeholder="Type your message here...")