Update handler.py

handler.py

@@ -1,28 +1,47 @@
+# handler.py
 
+import os
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
+# these globals will be filled in by init()
+generator = None
+
 def init():
     """
-    Called once at container startup. We load the fine-tuned T5 model here.
+    Called once when the container starts up.
+    Load your fine-tuned Flan-T5 model & tokenizer from the repo root.
     """
     global generator
-    model_dir = "."  # root of repo
+
+    model_dir = "."  # root of the repository
+    device = 0 if torch.cuda.is_available() else -1
+
+    # load tokenizer + model
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
-    model     = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
-    # build a text2text-generation pipeline on GPU
+    model     = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to(device if device >= 0 else "cpu")
+
+    # build a HF pipeline for text2text-generation
     generator = pipeline(
         "text2text-generation",
         model=model,
         tokenizer=tokenizer,
-        device=0
+        device=device   # GPU 0 or CPU (-1)
     )
 
 def run(payload: dict) -> list:
     """
-    Called on every HTTP request. 
-    Expects JSON: { "inputs": "<string>", "parameters": { ... } }
-    Returns the pipeline output, typically a list of {generated_text: ...}.
+    Called on every HTTP request.
+    Expects JSON with:
+      {
+        "inputs": "<your prompt string>",
+        "parameters": { ...generation kwargs… }
+      }
+    Returns a list-of-dicts, e.g.:
+      [ { "generated_text": "…" } ]
     """
-    text   = payload.get("inputs", "")
+    prompt = payload.get("inputs", "")
     params = payload.get("parameters", {})
-    return generator(text, **params)
+
+    # run the pipeline and return its output directly
+    return generator(prompt, **params)