Update app.py

app.py

@@ -11,8 +11,8 @@ SYSTEM_PROMPT = "You are a helpful and friendly AI assistant."
 # Log in using the secret token
 login(token=getenv("HF_TOKEN"))
 
-# Load Gemma 2B with optimizations for CPU
-model_name = "google/gemma-2b-270m"
+# Load the specified model with CPU optimizations
+model_name = "google/gemma-3-270m"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -21,32 +21,38 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="cpu"  # Explicitly map to CPU
 )
 
-# Simplify Gradio interface
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Gemma 2B Chatbot (CPU-Optimized)")
-    
+    gr.Markdown("# Gemma 3 270M Chatbot (CPU-Optimized)")
+
     with gr.Row():
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Chat")
+            chatbot = gr.Chatbot(label="Gemma 3 Chat")
             text_input = gr.Textbox(label="Your message")
             submit_button = gr.Button("Send")
-        
+
         with gr.Column(scale=1):
-            gr.Markdown("## Settings")
-            max_length_slider = gr.Slider(minimum=20, maximum=512, value=100, label="Max New Tokens")
-            temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, label="Temperature")
+            gr.Markdown("## User Controls")
+            max_length_slider = gr.Slider(minimum=20, maximum=512, value=100, label="Max New Tokens")  # Reduced max for CPU
+            temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
+            top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p")
+            top_k_slider = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k")
 
     def build_gemma_prompt(chat_history, new_message):
-        # Simplified prompt in Gemma's expected format
-        prompt = f"{SYSTEM_PROMPT}\n\n"
-        for user_msg, model_msg in chat_history:
+        # Simplified prompt construction in Gemma format
+        prompt = ""
+        for i, (user_msg, model_msg) in enumerate(chat_history):
+            if i == 0:
+                user_msg = f"{SYSTEM_PROMPT}\n\n{user_msg}"
             prompt += f"<start_of_turn>user\n{user_msg}<end_of_turn>\n"
             if model_msg:
                 prompt += f"<start_of_turn>model\n{model_msg}<end_of_turn>\n"
+        if not chat_history:
+            new_message = f"{SYSTEM_PROMPT}\n\n{new_message}"
         prompt += f"<start_of_turn>user\n{new_message}<end_of_turn>\n<start_of_turn>model\n"
         return prompt
 
-    def respond(message, chat_history, max_length, temperature):
+    def respond(message, chat_history, max_length, temperature, top_p, top_k):
         # Build prompt
         full_prompt = build_gemma_prompt(chat_history, message)
         
@@ -56,14 +62,14 @@ with gr.Blocks() as demo:
         # Update UI history
         chat_history.append((message, ""))
         
-        # Initialize streamer with proper token skipping
+        # Initialize streamer with proper token handling
         streamer = TextIteratorStreamer(
             tokenizer,
             skip_prompt=True,
             skip_special_tokens=True,
-            clean_up_tokenization_spaces=True  # Avoid gibberish from token artifacts
+            clean_up_tokenization_spaces=True  # Prevent token artifacts
         )
-        
+
         # Generation parameters
         generation_kwargs = {
             "input_ids": inputs["input_ids"],
@@ -71,14 +77,16 @@ with gr.Blocks() as demo:
             "streamer": streamer,
             "max_new_tokens": int(max_length),
             "temperature": float(temperature),
+            "top_p": float(top_p),
+            "top_k": int(top_k),
             "do_sample": True
         }
         
-        # Run generation in a separate thread with no_grad
+        # Run generation with no_grad for memory efficiency
         with torch.no_grad():
             thread = Thread(target=model.generate, kwargs=generation_kwargs)
             thread.start()
-            
+
             # Stream response
             accumulated_text = ""
             for new_text in streamer:
@@ -88,7 +96,7 @@ with gr.Blocks() as demo:
 
     submit_button.click(
         respond,
-        [text_input, chatbot, max_length_slider, temperature_slider],
+        [text_input, chatbot, max_length_slider, temperature_slider, top_p_slider, top_k_slider],
         [text_input, chatbot]
     )