Update app.py

app.py

@@ -5,7 +5,6 @@ import os
 
 # 🔧 CPU Optimization Suite
 os.environ["OMP_NUM_THREADS"] = "4"
-os.environ["MKL_NUM_THREADS"] = "4"
 torch.set_num_threads(4)
 torch.manual_seed(42)
 
@@ -22,16 +21,31 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 # ✅ Add pad_token if missing (required for batched generation)
 if tokenizer.pad_token is None:
+    # First add special token to tokenizer
     tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-    tokenizer.pad_token = tokenizer.eos_token  # Fallback to EOS as pad
+    
+    # Then resize model embeddings to accommodate new token
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+        cache_dir=cache_dir
+    )
+    model.resize_token_embeddings(len(tokenizer))
+    
+    # Finally set pad_token
+    tokenizer.pad_token = tokenizer.eos_token or tokenizer.cls_token or '[PAD]'
+else:
+    # Load model normally if pad_token exists
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+        cache_dir=cache_dir
+    )
 
-# 🧠 Load Model with CPU-specific settings
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float32,
-    low_cpu_mem_usage=True,
-    cache_dir=cache_dir
-).eval()
+# 🧠 Final model setup
+model = model.eval()
 
 def generate_response(prompt, max_new_tokens=128, temperature=0.7, top_p=0.9, num_sequences=1):
     """Optimized for 18GB CPU with strict memory control"""
@@ -54,7 +68,7 @@ def generate_response(prompt, max_new_tokens=128, temperature=0.7, top_p=0.9, nu
                 top_p=float(top_p),
                 do_sample=True,
                 num_return_sequences=int(num_sequences),
-                pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+                pad_token_id=tokenizer.convert_tokens_to_ids(tokenizer.pad_token),
                 eos_token_id=tokenizer.eos_token_id
             )