handling error

app.py

@@ -23,6 +23,21 @@ 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!
 """)
 
+# Check for Hugging Face Token
+huggingface_token = os.getenv("HF_TOKEN")
+if not huggingface_token:
+    st.warning("""
+    ⚠️ **No Hugging Face API token detected**
+    
+    The Gemma models require accepting a license and authentication to use.
+    
+    To make this app work:
+    1. Create a Hugging Face account
+    2. Accept the model license at: https://huggingface.co/google/gemma-2-2b-it
+    3. Create a HF token at: https://huggingface.co/settings/tokens
+    4. Add your token as a secret named 'HF_TOKEN' in your Space settings
+    """)
+
 # Sidebar with information
 with st.sidebar:
     st.header("About Gemma")
@@ -66,6 +81,8 @@ if 'generation_complete' not in st.session_state:
     st.session_state.generation_complete = False
 if 'generated_text' not in st.session_state:
     st.session_state.generated_text = ""
+if 'error_message' not in st.session_state:
+    st.session_state.error_message = None
 
 # Model parameters
 col1, col2 = st.columns(2)
@@ -83,110 +100,157 @@ user_input = st.text_area("Enter your prompt:",
                           placeholder="e.g., Write a short story about a robot discovering emotions")
 
 # Function to load model and generate text
-@st.cache_resource
+@st.cache_resource(show_spinner=False)
 def load_model():
-    # Get API Token
-    huggingface_token = os.getenv("HF_TOKEN")
-    if not huggingface_token:
-        st.warning("No Hugging Face API token found. Some models may not be accessible.")
-    
-    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
+    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):
-    tokenizer, model = load_model()
-    
-    # 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
-    progress_bar = st.progress(0)
-    status_text = st.empty()
-    output_area = st.empty()
-    
-    tokens_generated = 0
-    generated_text = ""
-    
-    # Generate with streaming
-    streamer_output = ""
-    
-    # Generate with step-by-step tracking for the progress bar
-    generate_kwargs = dict(
-        inputs=inputs["input_ids"],
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    
-    status_text.text("Generating response...")
-    
-    with torch.no_grad():
-        # Generate text step by step
-        for i in range(max_new_tokens):
-            if i == 0:
-                outputs = model.generate(
-                    **generate_kwargs,
-                    max_new_tokens=1,
-                )
-                generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
-            else:
-                input_ids = torch.cat([inputs["input_ids"], generated_ids], dim=1)
-                outputs = model.generate(
-                    input_ids=input_ids,
-                    max_new_tokens=1,
-                    do_sample=True,
-                    temperature=temperature,
-                    pad_token_id=tokenizer.eos_token_id
-                )
-                new_token = outputs[0][-1].unsqueeze(0)
-                generated_ids = torch.cat([generated_ids, new_token], dim=0)
-            
-            # Decode text
-            current_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
-            
-            # Update streaming output
-            streamer_output = current_text
-            
-            # 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 end token
-            if generated_ids[-1].item() == tokenizer.eos_token_id:
-                break
+    try:
+        with st.spinner("Loading model... (this may take a minute on first run)"):
+            tokenizer, model = load_model()
+        
+        # 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
+        progress_bar = st.progress(0)
+        status_text = st.empty()
+        output_area = st.empty()
+        
+        tokens_generated = 0
+        generated_text = ""
+        
+        # Generate with streaming
+        streamer_output = ""
+        
+        # Generate with step-by-step tracking for the progress bar
+        generate_kwargs = dict(
+            inputs=inputs["input_ids"],
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        status_text.text("Generating response...")
+        
+        with torch.no_grad():
+            # Generate text step by step
+            for i in range(max_new_tokens):
+                if i == 0:
+                    outputs = model.generate(
+                        **generate_kwargs,
+                        max_new_tokens=1,
+                    )
+                    generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
+                else:
+                    input_ids = torch.cat([inputs["input_ids"], generated_ids], dim=1)
+                    outputs = model.generate(
+                        input_ids=input_ids,
+                        max_new_tokens=1,
+                        do_sample=True,
+                        temperature=temperature,
+                        pad_token_id=tokenizer.eos_token_id
+                    )
+                    new_token = outputs[0][-1].unsqueeze(0)
+                    generated_ids = torch.cat([generated_ids, new_token], dim=0)
+                
+                # Decode text
+                current_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+                
+                # Update streaming output
+                streamer_output = current_text
                 
-            # Add a small delay to simulate typing
-            time.sleep(0.01)
+                # 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 end token
+                if generated_ids[-1].item() == tokenizer.eos_token_id:
+                    break
+                    
+                # Add a small delay to simulate typing
+                time.sleep(0.01)
+        
+        status_text.text("Generation complete!")
+        progress_bar.progress(1.0)
+        
+        return streamer_output
     
-    status_text.text("Generation complete!")
-    progress_bar.progress(1.0)
+    except Exception as e:
+        st.session_state.error_message = str(e)
+        return None
+
+# Show any existing error
+if st.session_state.error_message:
+    st.error(f"Error: {st.session_state.error_message}")
     
-    return streamer_output
+    # Add troubleshooting information
+    with st.expander("Troubleshooting Information"):
+        st.markdown("""
+        ### Common Issues:
+        
+        1. **Missing Hugging Face Token**: The Gemma model requires authentication. Add your token as a secret named 'HF_TOKEN' in the Space settings.
+        
+        2. **License Acceptance**: You need to accept the model license on the [Gemma model page](https://huggingface.co/google/gemma-2-2b-it).
+        
+        3. **Internet Connection**: The model needs to be downloaded the first time the app runs. Ensure your Space has internet access.
+        
+        4. **Resource Constraints**: The Gemma model requires significant resources. Consider upgrading your Space's hardware if you're encountering memory issues.
+        
+        ### How to Fix:
+        
+        1. Create a [Hugging Face account](https://huggingface.co/join)
+        2. Visit the [Gemma model page](https://huggingface.co/google/gemma-2-2b-it) and accept the license
+        3. Create a token at https://huggingface.co/settings/tokens
+        4. Add your token to the Space: Settings → Secrets → New Secret (HF_TOKEN)
+        """)
 
 # Generate button
 if st.button("Generate Text"):
-    if user_input:
+    # Reset any previous errors
+    st.session_state.error_message = None
+    
+    if not huggingface_token:
+        st.error("Hugging Face token is required! Please add your token as described above.")
+    elif user_input:
         st.session_state.user_prompt = user_input
-        with st.spinner("Generating text..."):
-            st.session_state.generated_text = generate_text(user_input, max_length, temperature)
+        result = generate_text(user_input, max_length, temperature)
+        if result is not None:  # Only set if no error occurred
+            st.session_state.generated_text = result
             st.session_state.generation_complete = True
     else:
         st.error("Please enter a prompt first!")
 
 # Display results
-if st.session_state.generation_complete:
+if st.session_state.generation_complete and not st.session_state.error_message:
     st.markdown("### Generated Text")
     st.markdown(st.session_state.generated_text)
     
@@ -207,6 +271,6 @@ st.markdown("---")
 st.markdown("""
 <div style="text-align: center">
 <p>Created with ❤️ | Powered by Gemma 2-2B-IT and Hugging Face</p>
-<p>Code available on <a href="https://huggingface.co/spaces/your-username/GemmaTextAppeal">Hugging Face Spaces</a></p>
+<p>Code available on <a href="https://huggingface.co/spaces" target="_blank">Hugging Face Spaces</a></p>
 </div>
 """, unsafe_allow_html=True)

requirements.txt

@@ -1,5 +1,6 @@
+
 streamlit==1.24.0
 torch>=2.0.0
-transformers>=4.31.0
+transformers>=4.34.0
 python-dotenv==1.0.0
 accelerate>=0.20.0