Update app.py

app.py

@@ -1,10 +1,5 @@
 import os
-import numpy as np
-import torch
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Handle Hugging Face Spaces GPU
+# Handle Spaces GPU
 if os.environ.get("SPACES_ZERO_GPU") is not None:
     import spaces
 else:
@@ -18,6 +13,12 @@ else:
 @spaces.GPU
 def fake_gpu():
     pass
+    
+import numpy as np
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import spaces
 
 # Available models
 AVAILABLE_MODELS = {
@@ -45,6 +46,7 @@ def load_model(model_name):
 # Load the default model at startup
 load_model("distilgpt2")
 
+@spaces.GPU()
 def get_next_token_predictions(text, model_name, top_k=10):
     """Generate the next token predictions with their probabilities."""
     global current_model, current_tokenizer