Update app.py

app.py

@@ -2,59 +2,85 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import spaces
+import random
+import numpy as np
 
 app = FastAPI()
 
+# Configuration
 model_id = "raduqus/reco_1b_16bit"
+device = "cuda"
+MAX_SEED = np.iinfo(np.int32).max
 
-# Initialize ZeroGPU
-spaces.gpu()
-
-try:
-    # Use spaces.gpu() decorator for initialization
-    @spaces.GPU
-    def load_model():
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-            device_map="auto"
-        )
-        return tokenizer, model
-
-    tokenizer, model = load_model()
-    print("Model loaded successfully on ZeroGPU.")
-except Exception as e:
-    print(f"Model loading error: {e}")
+def infer(
+    prompt, 
+    negative_prompt=None, 
+    seed=0, 
+    randomize_seed=True, 
+    max_length=100, 
+    temperature=0.7, 
+    top_p=0.9
+):
+    # Seed handling
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    # Set random generator
+    generator = torch.Generator().manual_seed(seed)
+    
+    # Load model
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, 
+        torch_dtype=torch.float16, 
+        variant="fp16", 
+        use_safetensors=True
+    )
+    
+    # Move to GPU
+    model = model.to(device)
+    
+    # Generate
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(
+        inputs.input_ids, 
+        max_length=max_length,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+        generator=generator
+    )
+    
+    # Decode
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return generated_text
 
 class RecommendationRequest(BaseModel):
     prompt: str
+    negative_prompt: str = None
+    seed: int = 0
+    randomize_seed: bool = True
     max_length: int = 100
     temperature: float = 0.7
     top_p: float = 0.9
 
 @app.post("/recommend")
-@spaces.GPU  # Ensure GPU usage for inference
 async def recommend_task(request: RecommendationRequest):
     try:
-        inputs = tokenizer(request.prompt, return_tensors="pt")
-        outputs = model.generate(
-            inputs.input_ids.to('cuda'), 
+        result = infer(
+            prompt=request.prompt,
+            negative_prompt=request.negative_prompt,
+            seed=request.seed,
+            randomize_seed=request.randomize_seed,
             max_length=request.max_length,
             temperature=request.temperature,
-            top_p=request.top_p,
-            do_sample=True
+            top_p=request.top_p
         )
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return {"recommendation": generated_text}
+        return {"recommendation": result}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/")
-async def root():
-    return {"message": "Task recommender is running on ZeroGPU!"}
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)