Update app.py

app.py

@@ -1,88 +1,73 @@
 import os
+import sys
 import torch
 from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
 
-# Disable Triton & set proper environment variables for HF Spaces
+# Create a directory with proper permissions
+os.makedirs("/tmp/hf_cache", exist_ok=True)
+
+# Set environment variables
 os.environ["TRITON_DISABLE"] = "1"
-os.environ["BNB_DISABLE_TRITON"] = "1"
+os.environ["BNB_DISABLE_TRITON"] = "1" 
 os.environ["USE_TORCH"] = "1"
 os.environ["BITSANDBYTES_NOWELCOME"] = "1"
+os.environ["HF_HOME"] = "/tmp/hf_cache"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
+os.environ["TORCH_HOME"] = "/tmp/hf_cache"
 
-# Set writable cache locations (HF Spaces needs explicit cache dirs)
-HF_CACHE_DIR = "/app/.cache/huggingface"
-TORCH_CACHE_DIR = "/app/.cache/torch"
-
-os.environ["HF_HOME"] = HF_CACHE_DIR
-os.environ["TRANSFORMERS_CACHE"] = HF_CACHE_DIR
-os.environ["TORCH_HOME"] = TORCH_CACHE_DIR
-
-# Create necessary directories & fix permissions
-for cache_dir in [HF_CACHE_DIR, TORCH_CACHE_DIR, "/tmp"]:
-    os.makedirs(cache_dir, exist_ok=True)
-    os.chmod(cache_dir, 0o777)  # Ensure all users can read/write
-
-# Initialize FastAPI
+# FastAPI app
 app = FastAPI()
 
-# Load base model & tokenizer
+# Load the base model
 base_model_name = "unsloth/Llama-3.2-3B-Instruct"
-adapter_name = "Suguru1846/lora_model_counseling_4bit"
-
-print("🚀 Loading base model...")
-try:
-    tokenizer = AutoTokenizer.from_pretrained(
-        base_model_name, cache_dir=HF_CACHE_DIR, trust_remote_code=True
-    )
-    
-    model = AutoModelForCausalLM.from_pretrained(
-        base_model_name,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        cache_dir=HF_CACHE_DIR,
-        trust_remote_code=True
-    )
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, cache_dir="/tmp/hf_cache")
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    cache_dir="/tmp/hf_cache"
+)
 
-    print("✅ Base model loaded successfully")
+# Try different PEFT versions
+adapter_name = "Suguru1846/lora_model_counseling_4bit"
+model = base_model  # Default to base model
+adapter_loaded = False
 
-    # Try loading LoRA adapter
-    print("🔄 Attempting to load LoRA adapter...")
+peft_versions = ["0.3.0", "0.4.0", "0.5.0"]
+for version in peft_versions:
     try:
-        model = PeftModel.from_pretrained(
-            model,
-            adapter_name,
-            device_map="auto",
-            cache_dir=HF_CACHE_DIR
-        )
-        print("✅ LoRA adapter loaded successfully")
-    except Exception as adapter_error:
-        print(f"⚠️ LoRA adapter loading failed: {adapter_error}")
-        print("⚠️ Running with base model only")
-
-except Exception as model_error:
-    print(f"❌ Error loading base model: {model_error}")
-    print("🔄 Falling back to OPT-350M model...")
-
-    # Fallback model in case of failure
-    base_model_name = "facebook/opt-350m"
-    
-    tokenizer = AutoTokenizer.from_pretrained(
-        base_model_name, cache_dir=HF_CACHE_DIR
-    )
-
-    model = AutoModelForCausalLM.from_pretrained(
-        base_model_name,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        cache_dir=HF_CACHE_DIR
-    )
+        # Add PEFT version to path
+        version_path = f"/app/peft_versions/{version}"
+        if version_path not in sys.path:
+            sys.path.insert(0, version_path)
+        
+        # Import PEFT from this version
+        from peft import PeftModel
+        
+        # Try loading with this version
+        print(f"Trying PEFT version {version}")
+        adapter_model = PeftModel.from_pretrained(base_model, adapter_name)
+        model = adapter_model
+        adapter_loaded = True
+        print(f"Success! Adapter loaded with PEFT version {version}")
+        break
+    except Exception as e:
+        print(f"Failed with PEFT version {version}: {str(e)}")
+        # Remove this version from path
+        if version_path in sys.path:
+            sys.path.remove(version_path)
+        
+        # Reset imports to try next version
+        if "peft" in sys.modules:
+            del sys.modules["peft"]
 
-    print("✅ Using fallback OPT-350M model")
+if not adapter_loaded:
+    print("Could not load adapter with any PEFT version. Using base model only.")
 
 @app.post("/generate")
 async def generate_text(prompt: str, max_tokens: int = 50):
-    """Generates text using the loaded model."""
+    """Generates text using the Llama model."""
     try:
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         outputs = model.generate(
@@ -95,11 +80,10 @@ async def generate_text(prompt: str, max_tokens: int = 50):
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return {"response": response}
     except Exception as e:
-        print(f"❌ Error generating text: {str(e)}")
+        print(f"Error generating text: {str(e)}")
         return {"error": str(e)}
 
 @app.get("/")
 async def root():
-    """Health check endpoint."""
-    model_type = "Base Llama-3.2 with adapter" if hasattr(model, "peft_config") else "Base Llama-3.2 only"
-    return {"message": f"AI Model is Running! Using: {model_type}"}
+    model_status = "with adapter" if adapter_loaded else "base model only"
+    return {"message": f"Model is running ({model_status})"}