ZeroGPU 0.1

core/compressor.py

@@ -1,4 +1,5 @@
 import torch
+import spaces
 import config
 from core.tokenizer_utils import count_tokens
 from models.model_loader import get_llm
@@ -12,7 +13,7 @@ TEXT:
 
 COMPRESSED:"""
 
-
+@spaces.GPU
 def _generate(prompt: str) -> str:
     model, tokenizer = get_llm()
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

models/model_loader.py

@@ -43,7 +43,7 @@ def switch_llm(model_id: str) -> str:
 def _load_llm(model_id: str):
     """Load model + its paired tokenizer. Both come from the same model_id."""
     global _llm, _tokenizer, _current_model_id
-    _tokenizer = AutoTokenizer.from_pretrained(model_id)
+    _tokenizer = AutoTokenizer.from_pretrained(model_id).to("cuda")
     _llm = AutoModelForCausalLM.from_pretrained(
         model_id,
         torch_dtype=torch.float32,

requirements.txt

@@ -5,4 +5,5 @@ torch>=2.2.0
 numpy>=1.26.0
 pandas>=2.0.0
 accelerate>=0.30.0
-huggingface_hub<2.0
+huggingface_hub<2.0
+spaces==0.50.4