Resolving the configuration problem

app.py

@@ -1,10 +1,12 @@
 import streamlit as st
-from transformers import AutoTokenizer, StoppingCriteria, StoppingCriteriaList
+from transformers import AutoTokenizer
 from peft import AutoPeftModelForCausalLM
 import torch
 import re
+from transformers import StoppingCriteria, StoppingCriteriaList
 import os
 
+# Set cache directory and get token
 os.environ['HF_HOME'] = '/app/cache'
 hf_token = os.getenv('HF_TOKEN')
 
@@ -19,34 +21,65 @@ class StopWordCriteria(StoppingCriteria):
 
 def load_model():
     try:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # Ensure cache directory exists
+        cache_dir = '/app/cache'
+        os.makedirs(cache_dir, exist_ok=True)
+        
+        # Check for HF token
+        if not hf_token:
+            st.warning("HuggingFace token not found. Some models may not be accessible.")
+        
+        # Check CUDA availability
         if torch.cuda.is_available():
+            device = torch.device("cuda")
             st.success(f"Using GPU: {torch.cuda.get_device_name(0)}")
         else:
-            st.warning("Using CPU for inference")
+            device = torch.device("cpu")
+            st.warning("CUDA is not available. Using CPU.")
 
+        # Fine-tuned model for generating scripts
         model_name = "Sidharthan/gemma2_scripter"
         
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            token=hf_token
-        )
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                trust_remote_code=True,
+                token=hf_token,
+                cache_dir=cache_dir
+            )
+        except Exception as e:
+            st.error(f"Error loading tokenizer: {str(e)}")
+            if "401" in str(e):
+                st.error("Authentication error. Please check your HuggingFace token.")
+            raise e
         
-        model = AutoPeftModelForCausalLM.from_pretrained(
-            model_name,
-            device_map=None,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-            cache_dir='/app/cache'
+        try:
+            # Load model with appropriate device settings
+            model = AutoPeftModelForCausalLM.from_pretrained(
+                model_name,
+                device_map=None,  # We'll handle device placement manually
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                token=hf_token,
+                cache_dir=cache_dir
+            )
             
-        ).to(device)
-        
-        return model, tokenizer
+            # Move model to device
+            model = model.to(device)
+            
+            return model, tokenizer
+
+        except Exception as e:
+            st.error(f"Error loading model: {str(e)}")
+            if "401" in str(e):
+                st.error("Authentication error. Please check your HuggingFace token.")
+            elif "disk space" in str(e).lower():
+                st.error("Insufficient disk space in cache directory.")
+            raise e
 
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
+        st.error(f"General error during model loading: {str(e)}")
         raise e
 
 def generate_script(tags, model, tokenizer, params):
@@ -77,6 +110,8 @@ def generate_script(tags, model, tokenizer, params):
             stopping_criteria=stopping_criteria
         )
         
+        # Move outputs back to CPU for decoding
+        outputs = outputs.cpu()
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         # Clean up response
@@ -111,38 +146,43 @@ def main():
     def get_model():
         return load_model()
     
-    model, tokenizer = get_model()
-    
-    # Tag input section
-    st.markdown("### Add Tags")
-    st.markdown("Enter tags separated by commas to generate a YouTube script")
-    
-    # Create columns for tag input and generate button
-    col1, col2 = st.columns([3, 1])
-    
-    with col1:
-        tags = st.text_input("Enter tags", placeholder="tech, AI, future, innovations...")
-    
-    with col2:
-        generate_button = st.button("Generate Script", type="primary")
-    
-    # Generated script section
-    if generate_button and tags:
-        st.markdown("### Generated Script")
-        with st.spinner("Generating script..."):
-            script = generate_script(tags, model, tokenizer, params)
-            st.text_area("Your script:", value=script, height=400)
+    try:
+        model, tokenizer = get_model()
+        
+        # Tag input section
+        st.markdown("### Add Tags")
+        st.markdown("Enter tags separated by commas to generate a YouTube script")
+        
+        # Create columns for tag input and generate button
+        col1, col2 = st.columns([3, 1])
+        
+        with col1:
+            tags = st.text_input("Enter tags", placeholder="tech, AI, future, innovations...")
+        
+        with col2:
+            generate_button = st.button("Generate Script", type="primary")
+        
+        # Generated script section
+        if generate_button and tags:
+            st.markdown("### Generated Script")
+            with st.spinner("Generating script..."):
+                script = generate_script(tags, model, tokenizer, params)
+                st.text_area("Your script:", value=script, height=400)
+                
+                # Add download button
+                st.download_button(
+                    label="Download Script",
+                    data=script,
+                    file_name="youtube_script.txt",
+                    mime="text/plain"
+                )
+        
+        elif generate_button and not tags:
+            st.warning("Please enter some tags first!")
             
-            # Add download button
-            st.download_button(
-                label="Download Script",
-                data=script,
-                file_name="youtube_script.txt",
-                mime="text/plain"
-            )
-    
-    elif generate_button and not tags:
-        st.warning("Please enter some tags first!")
+    except Exception as e:
+        st.error("Failed to initialize the application. Please check the logs for details.")
+        st.error(f"Error: {str(e)}")
 
 if __name__ == "__main__":
     main()