Update app.py

app.py

@@ -1,16 +1,38 @@
+import os
+import torch
 from flask import Flask, request, jsonify
 import requests
-import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
+# Set cache directory explicitly
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface_cache'
+os.makedirs('/tmp/huggingface_cache', exist_ok=True)
+
 app = Flask(__name__)
 
 # Define the SearXNG instance URL
 SEARXNG_INSTANCE_URL = "https://oscarwang2-searxng.hf.space/search"
 
-# Load the educational content classifier
-tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/fineweb-edu-classifier")
-model = AutoModelForSequenceClassification.from_pretrained("HuggingFaceTB/fineweb-edu-classifier")
+# Load the educational content classifier with explicit cache directory
+def load_model_with_retry(max_retries=3):
+    for attempt in range(max_retries):
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(
+                "HuggingFaceTB/fineweb-edu-classifier", 
+                cache_dir='/tmp/huggingface_cache'
+            )
+            model = AutoModelForSequenceClassification.from_pretrained(
+                "HuggingFaceTB/fineweb-edu-classifier", 
+                cache_dir='/tmp/huggingface_cache'
+            )
+            return tokenizer, model
+        except Exception as e:
+            print(f"Model loading attempt {attempt + 1} failed: {e}")
+            if attempt == max_retries - 1:
+                raise
+
+# Load models at startup
+tokenizer, model = load_model_with_retry()
 
 def classify_educational_quality(text):
     """