Update app.py

app.py

@@ -19,60 +19,89 @@ from huggingface_hub import login
 
 app = Flask(__name__)
 
+# Configuration for Hugging Face Spaces
 PORT = int(os.environ.get("PORT", 7860))
 
-hf_token = os.environ.get("HF_TOKEN")
-if hf_token:
-    login(hf_token)
-
-UPLOAD_FOLDER = '/tmp/uploads'  # Change to tmp directory for Spaces
+# Set cache directories to /tmp
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
+os.environ['HF_HOME'] = '/tmp/hf_home'
+os.environ['XDG_CACHE_HOME'] = '/tmp/cache'
+os.environ['HF_DATASETS_CACHE'] = '/tmp/datasets_cache'
+
+# Create necessary directories with proper permissions
+for directory in [
+    '/tmp/transformers_cache',
+    '/tmp/hf_home',
+    '/tmp/cache',
+    '/tmp/datasets_cache',
+    '/tmp/uploads'
+]:
+    os.makedirs(directory, exist_ok=True)
+
+# Configure upload folder inside the space
+UPLOAD_FOLDER = '/tmp/uploads'
 ALLOWED_EXTENSIONS = {'py'}
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+
 # Database configuration
 DATABASE_PATH = '/tmp/chat_database.db'
-CACHE_DIR = "/tmp/huggingface_cache"
-MODEL_CACHE_DIR = "/tmp/model_cache"
-os.makedirs(CACHE_DIR, exist_ok=True)
-os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
-os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR
-os.environ['HF_HOME'] = CACHE_DIR
-os.environ['HF_DATASETS_CACHE'] = CACHE_DIR
-
-# Initialize LangChain with Ollama LLM
-if hf_token:
-    model_name = "mistralai/Mistral-7B-Instruct-v0.1"
-else:
-    # Fallback to a free, smaller model
-    model_name = "microsoft/phi-4"
-
-try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=CACHE_DIR)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        load_in_8bit=True,
-        cache_dir=MODEL_CACHE_DIR
-    )
 
-    # Create pipeline
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.95,
-        repetition_penalty=1.15
-    )
+def get_model_name():
+    """Determine which model to use based on token availability"""
+    try:
+        hf_token = os.environ.get("HF_TOKEN")
+        if hf_token:
+            # Set token in environment and return gated model name
+            os.environ['HUGGING_FACE_HUB_TOKEN'] = hf_token
+            return "mistralai/Mistral-7B-Instruct-v0.1"
+        else:
+            # Return free model if no token
+            return "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    except Exception as e:
+        print(f"Error accessing token: {e}")
+        return "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-    # Initialize LangChain with HuggingFacePipeline
-    llm = HuggingFacePipeline(pipeline=pipe)
+def initialize_model():
+    """Initialize the model with appropriate settings"""
+    try:
+        model_name = get_model_name()
+        print(f"Initializing model: {model_name}")
+
+        # Initialize tokenizer with explicit cache directory
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            cache_dir='/tmp/transformers_cache',
+            token=os.environ.get('HUGGING_FACE_HUB_TOKEN')
+        )
+
+        # Initialize model with explicit cache directory
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            cache_dir='/tmp/transformers_cache',
+            token=os.environ.get('HUGGING_FACE_HUB_TOKEN'),
+            torch_dtype=torch.float16,
+            device_map="auto",
+            load_in_8bit=True
+        )
+
+        # Create pipeline
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.95,
+            repetition_penalty=1.15
+        )
+
+        return HuggingFacePipeline(pipeline=pipe)
+    except Exception as e:
+        print(f"Error initializing model: {e}")
+        raise
 
-except Exception as e:
-    print(f"Error loading model: {e}")
-    raise
+# Initialize LLM
+llm = initialize_model()
 
 @contextmanager
 def get_db_connection():