Update app.py

app.py

@@ -1,34 +1,14 @@
-import subprocess
-import sys
 import os
-
-# --- THE STABILIZER BLOCK ---
-print("🛠️ Stabilizing environment...")
-subprocess.check_call([
-    sys.executable, "-m", "pip", "install", 
-    "tokenizers==0.20.1",
-    "transformers==4.45.2",
-    "huggingface-hub==0.24.7",
-    "peft==0.13.2"
-])
-
 import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import PeftModel
-from huggingface_hub import login
 
-# Login with HuggingFace token
+# Get HuggingFace token from environment
 HF_TOKEN = os.environ.get("HF_TOKEN")
-if HF_TOKEN:
-    print("🔐 Logging in to HuggingFace...")
-    login(token=HF_TOKEN)
-else:
-    print("⚠️ No HF_TOKEN found - may fail on gated models")
 
 BASE_MODEL = "polyglots/SinLlama_v01"
 LORA_ADAPTER = "E-motionAssistant/SinLlama_v01-Therapy-Sinhala"
-
 SYSTEM_PROMPT = "You are an empathetic Sinhala therapist providing mental health support."
 
 model = None
@@ -37,27 +17,35 @@ tokenizer = None
 def load_model():
     global model, tokenizer
     if model is None:
+        print("🔐 Loading with 4-bit quantization...")
+        
+        # 4-bit quantization config
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16
+        )
+        
         print(f"📥 Loading base model: {BASE_MODEL}...")
         base_model = AutoModelForCausalLM.from_pretrained(
             BASE_MODEL,
-            torch_dtype=torch.float32,
-            device_map="cpu",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-            ignore_mismatched_sizes=True,
-            token=HF_TOKEN
+            quantization_config=bnb_config,
+            device_map="auto",
+            token=HF_TOKEN,
+            trust_remote_code=True
         )
         
         print(f"📥 Loading LoRA adapter: {LORA_ADAPTER}...")
         model = PeftModel.from_pretrained(base_model, LORA_ADAPTER, token=HF_TOKEN)
         
-        print(f"📥 Loading tokenizer from adapter...")
-        tokenizer = AutoTokenizer.from_pretrained(LORA_ADAPTER, trust_remote_code=True, token=HF_TOKEN)
+        print(f"📥 Loading tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(LORA_ADAPTER, token=HF_TOKEN, trust_remote_code=True)
         
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
             
-        print("✅ Success: Sinhala Therapy System is online on CPU!")
+        print("✅ Model loaded in 4-bit!")
 
 load_model()
 
@@ -66,7 +54,6 @@ def chat(message, history):
         return ""
     
     try:
-        # Build prompt
         prompt = f"{SYSTEM_PROMPT}\n\n"
         
         for user_msg, bot_msg in history[-3:]:
@@ -74,7 +61,6 @@ def chat(message, history):
         
         prompt += f"User: {message}\nTherapist:"
         
-        # Tokenize
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
         
         with torch.no_grad():
@@ -88,15 +74,14 @@ def chat(message, history):
                 eos_token_id=tokenizer.eos_token_id
             )
         
-        # Decode only the new tokens
         input_len = inputs.input_ids.shape[1]
         response = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)
         
         return response.strip()
     
     except Exception as e:
-        print(f"❌ Generation Error: {e}")
-        return f"සමාවන්න, දෝෂයක් ඇතිවිය: {str(e)}. කරුණාකර නැවත උත්සාහ කරන්න."
+        print(f"❌ Error: {e}")
+        return f"සමාවන්න, දෝෂයක් ඇතිවිය. කරුණාකර නැවත උත්සාහ කරන්න."
 
 demo = gr.ChatInterface(
     fn=chat,