print of response

app.py

@@ -23,6 +23,43 @@ This app demonstrates the text generation capabilities of Google's Gemma 2-2B-IT
 Enter a prompt below and see the model generate text in real-time!
 """)
 
+# Function to load model
+@st.cache_resource(show_spinner=False)
+def load_model():
+    try:
+        # Get API Token
+        huggingface_token = os.getenv("HF_TOKEN")
+        if not huggingface_token:
+            return None, None, "No Hugging Face API token found. Please add your token as a secret named 'HF_TOKEN'."
+        
+        # Attempt to download model with explicit token
+        tokenizer = AutoTokenizer.from_pretrained(
+            "google/gemma-2-2b-it", 
+            token=huggingface_token
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            "google/gemma-2-2b-it",
+            token=huggingface_token,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        return tokenizer, model, None
+    except Exception as e:
+        return None, None, str(e)
+
+# Try to load the model at startup
+with st.spinner("Initializing the Gemma model... this may take a minute."):
+    tokenizer, model, load_error = load_model()
+
+if load_error:
+    st.error(f"Error loading model: {load_error}")
+else:
+    if tokenizer and model:
+        st.success("✅ Gemma model loaded successfully! Ready to generate text.")
+    else:
+        st.warning("⚠️ Model not loaded. Please check your Hugging Face token.")
+
 # Check for Hugging Face Token
 huggingface_token = os.getenv("HF_TOKEN")
 if not huggingface_token:
@@ -99,39 +136,12 @@ user_input = st.text_area("Enter your prompt:",
                           height=100, 
                           placeholder="e.g., Write a short story about a robot discovering emotions")
 
-# Function to load model and generate text
-@st.cache_resource(show_spinner=False)
-def load_model():
-    try:
-        # Get API Token
-        huggingface_token = os.getenv("HF_TOKEN")
-        if not huggingface_token:
-            raise ValueError("No Hugging Face API token found. Please add your token as a secret named 'HF_TOKEN'.")
-        
-        # Attempt to download model with explicit token
-        tokenizer = AutoTokenizer.from_pretrained(
-            "google/gemma-2-2b-it", 
-            token=huggingface_token,
-            use_fast=True
-        )
-        
-        model = AutoModelForCausalLM.from_pretrained(
-            "google/gemma-2-2b-it",
-            token=huggingface_token,
-            torch_dtype=torch.float16,
-            device_map="auto"
-        )
-        return tokenizer, model
-    except Exception as e:
-        # Re-raise the exception to be handled in the calling function
-        raise e
-
 def generate_text(prompt, max_new_tokens=300, temperature=0.7):
+    if not tokenizer or not model:
+        st.session_state.error_message = "Model not properly loaded. Please check your Hugging Face token."
+        return None
+
     try:
-        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"
         
@@ -141,40 +151,40 @@ def generate_text(prompt, max_new_tokens=300, temperature=0.7):
         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)
+        # Tokenize the input
+        encoding = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+        input_ids = encoding["input_ids"]
         
-        # Create an attention mask of ones (attend to all tokens)
+        # Ensure we have a proper attention mask
         attention_mask = torch.ones_like(input_ids)
         
-        # Generate the full text at once (simpler and more reliable)
-        generated_ids = model.generate(
+        # Simple approach - generate all at once
+        output_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,
+            pad_token_id=tokenizer.eos_token_id
         )
         
-        # Get only the newly generated tokens (exclude input prompt)
-        generated_text = tokenizer.decode(generated_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+        st.write("Generation completed, processing output...")
         
-        # Simulate token-by-token generation for visual effect
+        # Get only the generated part (exclude the prompt)
+        new_tokens = output_ids[0][input_ids.shape[1]:]
+        generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
+        
+        # Display incrementally for visual effect
+        display_text = ""
         words = generated_text.split()
-        displayed_text = ""
+        total_words = len(words)
         
         for i, word in enumerate(words):
-            displayed_text += word + " "
-            
-            # Update progress and display
-            progress = min(1.0, (i + 1) / len(words))
+            display_text += word + " "
+            progress = min(1.0, (i + 1) / total_words)
             progress_bar.progress(progress)
-            output_area.markdown(f"**Generated Response:**\n\n{displayed_text}")
-            
-            # Small delay for visual effect
-            time.sleep(0.05)
+            output_area.markdown(f"**Generated Response:**\n\n{display_text}")
+            time.sleep(0.05)  # Brief delay for visual effect
         
         status_text.text("Generation complete!")
         progress_bar.progress(1.0)
@@ -182,7 +192,8 @@ def generate_text(prompt, max_new_tokens=300, temperature=0.7):
         return generated_text
     
     except Exception as e:
-        st.session_state.error_message = str(e)
+        st.session_state.error_message = f"Error during generation: {str(e)}"
+        st.error(f"Error during generation: {str(e)}")
         return None
 
 # Show any existing error
@@ -210,6 +221,17 @@ if st.session_state.error_message:
         4. Add your token to the Space: Settings → Secrets → New Secret (HF_TOKEN)
         """)
 
+# Add a debug section
+with st.expander("Debug Information"):
+    st.write(f"Model loaded: {model is not None}")
+    st.write(f"Tokenizer loaded: {tokenizer is not None}")
+    st.write(f"Device mapping: {model.device_map if model else 'N/A'}")
+    st.write(f"Hugging Face token set: {huggingface_token is not None}")
+    if torch.cuda.is_available():
+        st.write(f"CUDA available: True (Device count: {torch.cuda.device_count()})")
+    else:
+        st.write("CUDA available: False")
+
 # Generate button
 if st.button("Generate Text"):
     # Reset any previous errors
@@ -219,7 +241,9 @@ if st.button("Generate Text"):
         st.error("Hugging Face token is required! Please add your token as described above.")
     elif user_input:
         st.session_state.user_prompt = user_input
+        st.write("Starting text generation...")
         result = generate_text(user_input, max_length, temperature)
+        st.write(f"Generation result: {'Success' if result else 'Failed'}")
         if result is not None:  # Only set if no error occurred
             st.session_state.generated_text = result
             st.session_state.generation_complete = True