Update main.py

main.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from typing import List, Optional, Dict, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer
 import torch
+import os
 
 app = FastAPI()
 
@@ -15,10 +16,11 @@ class PromptRequest(BaseModel):
 @app.on_event("startup")
 def load_model():
     global model, tokenizer, pipe
+    os.environ["TRANSFORMERS_CACHE"] = "./cache"
     model_path = "model/models--meta-llama--Llama-3.2-3B-Instruct/snapshots/0cb88a4f764b7a12671c53f0838cd831a0843b95"
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     streamer = TextStreamer(tokenizer=tokenizer, skip_prompt=True)
-    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, cache_dir="./cache")
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, streamer=streamer)
 
 @app.post("/generate/")