Refactor text generation logic to use Hugging Face API and remove local model loading

main.py

@@ -1,8 +1,13 @@
 import os
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline
 from huggingface_hub import login
+import requests
+
+API_URL = "https://router.huggingface.co/v1/chat/completions"
+headers = {
+    "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
+}
 
 # Login if token is available (e.g., via Hugging Face Spaces secrets)
 if "HF_TOKEN" in os.environ:
@@ -13,14 +18,9 @@ app = FastAPI()
 class Prompt(BaseModel):
     text: str
 
-@app.on_event("startup")
-async def load_model():
-    global generator
-    try:
-        generator = pipeline("text-generation", model="nvidia/gpt-oss-120b-Eagle3", use_auth_token=True)
-    except Exception:
-        # fallback to public model
-        generator = pipeline("text-generation", model="distilgpt2")
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()
 
 @app.get("/")
 def read_root():
@@ -28,5 +28,13 @@ def read_root():
 
 @app.post("/generate")
 async def generate_text(prompt: Prompt):
-    output = generator(prompt.text, max_new_tokens=200, do_sample=True, temperature=0.7)
-    return {"response": output[0]["generated_text"]}
+    response = query({
+        "messages": [
+            {
+                "role": "user",
+                "content": prompt.text
+            }
+        ],
+        "model": "meta-llama/Llama-3.2-3B-Instruct"
+    })
+    return {"response": response["choices"][0]["message"]}