Create service.py

service.py

@@ -0,0 +1,73 @@
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Global variables
+_model = None
+_tokenizer = None
+_model_name = "microsoft/DialoGPT-small"  # Using a smaller, faster model for testing
+
+def initialize_tokenizer():
+    """Initialize tokenizer"""
+    global _tokenizer
+    if _tokenizer is None:
+        print("[MinimalService] Loading tokenizer...")
+        _tokenizer = AutoTokenizer.from_pretrained(_model_name)
+        if _tokenizer.pad_token is None:
+            _tokenizer.pad_token = _tokenizer.eos_token
+        print("[MinimalService] Tokenizer loaded successfully.")
+    return _tokenizer
+
+@spaces.GPU
+def generate_text_gpu(prompt: str, max_tokens: int = 50):
+    """GPU function for text generation"""
+    global _model, _tokenizer
+    
+    print("[MinimalService] GPU function called")
+    
+    # Initialize tokenizer
+    if _tokenizer is None:
+        initialize_tokenizer()
+    
+    # Load model in GPU context
+    if _model is None:
+        print("[MinimalService] Loading model...")
+        _model = AutoModelForCausalLM.from_pretrained(
+            _model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        print("[MinimalService] Model loaded.")
+    
+    # Simple generation
+    inputs = _tokenizer.encode(prompt, return_tensors="pt")
+    device = next(_model.parameters()).device
+    inputs = inputs.to(device)
+    
+    with torch.no_grad():
+        outputs = _model.generate(
+            inputs,
+            max_new_tokens=max_tokens,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=_tokenizer.eos_token_id
+        )
+    
+    response = _tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+class MinimalService:
+    def __init__(self):
+        print("[MinimalService] Service initialized")
+        # Initialize tokenizer immediately
+        initialize_tokenizer()
+    
+    def generate(self, prompt: str):
+        """Public method to generate text"""
+        return generate_text_gpu(prompt)
+
+# Create instance
+service = MinimalService()
+
+# Print confirmation that GPU function is registered
+print(f"[MinimalService] GPU function available: {generate_text_gpu.__name__}")