Refactor main.py to implement model loading on startup and add Hugging Face token login

main.py

@@ -1,15 +1,27 @@
+import os
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
+from huggingface_hub import login
 
-# Load a small model to test (can change to bigger one later)
-generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
+# Login if token is available (e.g., via Hugging Face Spaces secrets)
+if "HF_TOKEN" in os.environ:
+    login(token=os.environ["HF_TOKEN"])
 
 app = FastAPI()
 
 class Prompt(BaseModel):
     text: str
 
+@app.on_event("startup")
+async def load_model():
+    global generator
+    try:
+        generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", use_auth_token=True)
+    except Exception:
+        # fallback to public model
+        generator = pipeline("text-generation", model="distilgpt2")
+
 @app.get("/")
 def read_root():
     return {"Hello": "docker embeddings!"}