Loads model before starting server

main.py

@@ -14,7 +14,8 @@ from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import uvicorn
 import torch
-
+import sys
+# torch.mps.empty_cache()
 
 app = FastAPI()
 
@@ -35,13 +36,29 @@ def format_prompt(system, message, history):
     prompt += {"role": "user", "content": message}
     return prompt
 
-def generate(item: Item):
+def setup():
     device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # if torch.backends.mps.is_available():
+    #     device = torch.device("mps")
+    #     x = torch.ones(1, device=device)
+    #     print (x)
+    # else:
+    #     device="cpu"
+    #     print ("MPS device not found.")
+    
+    # device = "auto"
+    # device=torch.device("cpu")
+    
     model_path = "ibm-granite/granite-34b-code-instruct-8k"
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     # drop device_map if running on CPU
     model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
     model.eval()
+    
+    return model, tokenizer, device
+
+def generate(item: Item, model, tokenizer, device):
     # change input text as desired
     chat = format_prompt(item.system_prompt, item.prompt, item.history)
     chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
@@ -56,9 +73,12 @@ def generate(item: Item):
     return output_text
 
 
+model, tokenizer, device = setup()
+
 @app.post("/generate/")
 async def generate_text(item: Item):
-    return {"response": generate(item)}
+    print(item, file=sys.stderr)
+    return {"response": generate(item, model, tokenizer, device)}
 
 @app.get("/")
 async def generate_text_root(item: Item):

requirements.txt

@@ -1,5 +1,5 @@
 fastapi
-uvicorn
+uvicorn[standard]
 huggingface_hub
 pydantic
 transformers