Update app.py

app.py

@@ -14,8 +14,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import PyPDF2
 import traceback
 import os
-import spaces
-from spaces import GPU
+
 
 import shutil
 from pathlib import Path
@@ -52,11 +51,11 @@ model = None
 tokenizer = None
 generation_config = None
 
-@spaces.GPU
+
 def test_llm_generation():
     try:
         test_prompt = "Hello, how are you today?"
-        inputs = tokenizer(test_prompt, return_tensors="pt")#.to(model.device)
+        inputs = tokenizer(test_prompt, return_tensors="pt").to(model.device)
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
@@ -70,7 +69,7 @@ def test_llm_generation():
     except Exception as e:
         add_log(f"❌ LLM quick test failed: {e}")
 
-@spaces.GPU
+
 def initialize_model():
     global model, tokenizer, generation_config
 
@@ -90,20 +89,20 @@ def initialize_model():
             add_log("✅ Set pad_token to eos_token")
 
         # Force GPU settings
-       # model = AutoModelForCausalLM.from_pretrained(
-       #     MODEL_ID,
-       #     torch_dtype=torch.float16,
-       #     cache_dir="/data/models",
-       #     trust_remote_code=True,
-       #     token=glotoken,
-       #     device_map={"": 0},  # <- force GPU:0
-       #     low_cpu_mem_usage=True
-       # )
         model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,  
+            MODEL_ID,
+            torch_dtype=torch.float16,
             cache_dir="/data/models",
-            trust_remote_code=True
+            trust_remote_code=True,
+            token=glotoken,
+            device_map={"": 0},  # <- force GPU:0
+            low_cpu_mem_usage=True
         )
+       # model = AutoModelForCausalLM.from_pretrained(
+       #     MODEL_ID,  
+       #     cache_dir="/data/models",
+       #     trust_remote_code=True
+       # )
         model.eval()
 
         generation_config = GenerationConfig(
@@ -198,8 +197,8 @@ Now format the following:
             truncation=True,
             max_length=2048
         )
-        #inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
-        inputs = {k: v for k, v in inputs.items()}
+        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+        #inputs = {k: v for k, v in inputs.items()}
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,