fix mask issue

app.py

@@ -131,74 +131,55 @@ def generate_text(prompt, max_new_tokens=300, temperature=0.7):
         with st.spinner("Loading model... (this may take a minute on first run)"):
             tokenizer, model = load_model()
         
+        # Simpler approach: use the model's built-in text generation capabilities
         # Format the prompt according to Gemma's expected format
         formatted_prompt = f"<bos><start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
         
-        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
-        
-        # Create the progress bar
+        # Create the progress bar and status indicators
         progress_bar = st.progress(0)
         status_text = st.empty()
         output_area = st.empty()
+        status_text.text("Generating response...")
+        
+        # Tokenize the input with attention mask explicitly set
+        encoding = tokenizer(formatted_prompt, return_tensors="pt")
+        input_ids = encoding["input_ids"].to(model.device)
+        
+        # Create an attention mask of ones (attend to all tokens)
+        attention_mask = torch.ones_like(input_ids)
+        
+        # Generate the full text at once (simpler and more reliable)
+        generated_ids = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            pad_token_id=tokenizer.eos_token_id,
+        )
         
-        streamer_output = ""
+        # Get only the newly generated tokens (exclude input prompt)
+        generated_text = tokenizer.decode(generated_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
         
-        status_text.text("Generating response...")
+        # Simulate token-by-token generation for visual effect
+        words = generated_text.split()
+        displayed_text = ""
         
-        with torch.no_grad():
-            # Generate text step by step to show progress
-            # Start with 1 token generation
-            input_ids = inputs["input_ids"]
-            generated_ids = None
+        for i, word in enumerate(words):
+            displayed_text += word + " "
+            
+            # Update progress and display
+            progress = min(1.0, (i + 1) / len(words))
+            progress_bar.progress(progress)
+            output_area.markdown(f"**Generated Response:**\n\n{displayed_text}")
             
-            for i in range(max_new_tokens):
-                if i == 0:
-                    # First token generation
-                    outputs = model.generate(
-                        input_ids=input_ids,
-                        max_new_tokens=1,
-                        do_sample=True,
-                        temperature=temperature,
-                        pad_token_id=tokenizer.eos_token_id
-                    )
-                    # Extract only the newly generated token(s)
-                    generated_ids = outputs[0][input_ids.shape[1]:].unsqueeze(0)
-                else:
-                    # For subsequent tokens, concatenate previous results
-                    current_input_ids = torch.cat([input_ids, generated_ids], dim=1)
-                    outputs = model.generate(
-                        input_ids=current_input_ids,
-                        max_new_tokens=1,
-                        do_sample=True,
-                        temperature=temperature,
-                        pad_token_id=tokenizer.eos_token_id
-                    )
-                    # Extract only the newly generated token
-                    new_token = outputs[0][-1].unsqueeze(0).unsqueeze(0)
-                    generated_ids = torch.cat([generated_ids, new_token], dim=1)
-                
-                # Decode the current state
-                current_output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-                streamer_output = current_output
-                
-                # Update progress and output
-                progress = min(1.0, (i + 1) / max_new_tokens)
-                progress_bar.progress(progress)
-                
-                # Update display
-                output_area.markdown(f"**Generated Response:**\n\n{streamer_output}")
-                
-                # Check if we've reached an EOS token in the latest output
-                if outputs[0][-1].item() == tokenizer.eos_token_id:
-                    break
-                
-                # Add a small delay to simulate typing
-                time.sleep(0.01)
+            # Small delay for visual effect
+            time.sleep(0.05)
         
         status_text.text("Generation complete!")
         progress_bar.progress(1.0)
         
-        return streamer_output
+        return generated_text
     
     except Exception as e:
         st.session_state.error_message = str(e)