change code + readme

README.md

@@ -1,8 +1,8 @@
 ---
 title: GemmaTextAppeal
-emoji: 👁
+emoji: 🩵
 colorFrom: purple
-colorTo: gray
+colorTo: blue
 sdk: streamlit
 sdk_version: 1.44.1
 app_file: app.py

app.py

@@ -38,12 +38,25 @@ def load_model():
             token=huggingface_token
         )
         
+        # Load model - use CPU configuration if no GPU available
+        model_kwargs = {
+            "token": huggingface_token,
+            "torch_dtype": torch.float16
+        }
+        
+        # Only add device_map if GPU is available
+        if torch.cuda.is_available():
+            model_kwargs["device_map"] = "auto"
+        
         model = AutoModelForCausalLM.from_pretrained(
             "google/gemma-2-2b-it",
-            token=huggingface_token,
-            torch_dtype=torch.float16,
-            device_map="auto"
+            **model_kwargs
         )
+        
+        # Move model to CPU if no GPU
+        if not torch.cuda.is_available():
+            model = model.to("cpu")
+            
         return tokenizer, model, None
     except Exception as e:
         return None, None, str(e)
@@ -168,8 +181,6 @@ def generate_text(prompt, max_new_tokens=300, temperature=0.7):
             pad_token_id=tokenizer.eos_token_id
         )
         
-        st.write("Generation completed, processing output...")
-        
         # 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)
@@ -225,7 +236,7 @@ if st.session_state.error_message:
 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"Device: {model.device 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()})")
@@ -241,12 +252,11 @@ 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
+        with st.spinner("Generating text..."):
+            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!")
 

requirements.txt

@@ -1,6 +1,7 @@
-
-streamlit==1.24.0
-torch>=2.0.0
-transformers>=4.34.0
+streamlit==1.30.0
+torch==2.1.0
+transformers==4.35.0
 python-dotenv==1.0.0
-accelerate>=0.20.0
+huggingface-hub==0.19.0
+accelerate==0.23.0
+protobuf==3.20.3