feat: huh

main.py

@@ -114,7 +114,35 @@ def load_model():
         logger.info("🧠 Loading ONNX decoder...")
         decoder_model_path = os.path.join(model_dir, "decoder_model_merged_q4.onnx")
         
-       decoder_session is None or tokenizer is None:
+        providers = ["CPUExecutionProvider"]
+        try:
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        except:
+            pass
+        
+        decoder_session = ort.InferenceSession(decoder_model_path, providers=providers)
+        logger.info(f"✅ Model loaded successfully (using {decoder_session.get_providers()})")
+
+    except Exception as e:
+        logger.exception("❌ Failed to load model")
+        raise e
+
+
+# =========================
+# 📥 REQUEST MODEL
+# =========================
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_tokens: Optional[int] = 100
+    temperature: Optional[float] = 0.7
+
+
+# =========================
+# 📤 RESPONSE ENDPOINT
+# =========================
+@app.post("/generate")
+async def generate(req: GenerateRequest):
+    if decoder_session is None or tokenizer is None:
         raise HTTPException(status_code=500, detail="Model not loaded")
 
     try:
@@ -182,7 +210,7 @@ def load_model():
             # Update past_key_values from outputs
             for j, output in enumerate(decoder_session.get_outputs()):
                 if output.name.startswith("past_key"):
-                    if jdecoder_sessionen(outputs):
+                    if j < len(outputs):
                         past_key_values[output.name] = outputs[j]
             
             # Check for EOS token
@@ -194,34 +222,6 @@ def load_model():
         
         logger.info("✅ Generation successful")
         
-
-# =========================
-# 📤 RESPONSE ENDPOINT
-# =========================
-@app.post("/generate")
-async def generate(req: GenerateRequest):
-    if model is None or tokenizer is None:
-        raise HTTPException(status_code=500, detail="Model not loaded")
-
-    try:
-        logger.info(f"🧠 Generating for prompt: {req.prompt[:50]}...")
-
-        # Tokenize input
-        inputs = tokenizer(req.prompt, return_tensors="pt")
-        
-        # Generate
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=req.max_tokens,
-            temperature=req.temperature,
-            do_sample=True,
-        )
-        
-        # Decode output
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        logger.info("✅ Generation successful")
-
         return {
             "success": True,
             "response": result
@@ -240,7 +240,7 @@ async def generate(req: GenerateRequest):
 async def health():
     return {
         "status": "ok",
-        "model_loaded": pipe is not None
+        "model_loaded": decoder_session is not None
     }