added more speed

main.py

@@ -1,33 +1,45 @@
 import os
 import torch
+import multiprocessing
 from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from dotenv import load_dotenv
+from accelerate import Accelerator
 
 # Load environment variables from a .env file (useful for local development)
 load_dotenv()
 
 # Initialize FastAPI app
-app = FastAPI(description="Use the Llama-3.2-1B-Instruct model using the api !!",docs_url="/",redoc_url="/doc")
+app = FastAPI(description="Use the Llama-3.2-1B-Instruct model using the API", docs_url="/", redoc_url="/doc")
 
 # Set your Hugging Face token from environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 MODEL = "meta-llama/Llama-3.2-1B-Instruct"
 
+# Auto-select CPU or GPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
+# Set PyTorch to use all available CPU cores if running on CPU
+torch.set_num_threads(multiprocessing.cpu_count())
+
+# Initialize Accelerator for managing device allocation
+accelerator = Accelerator()
+
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     token=HF_TOKEN,
-    torch_dtype=torch.float16,  # Use float16 for better GPU memory usage
+    torch_dtype=torch.bfloat16 if device == 'cpu' else torch.float16,  # Use bfloat16 for CPUs, float16 for GPUs
     device_map="auto"
 )
 
+# Prepare model for multi-device setup with accelerate
+model, tokenizer = accelerator.prepare(model, tokenizer)
+
 # Pydantic model for input
 class PromptRequest(BaseModel):
     prompt: str
@@ -49,4 +61,3 @@ async def generate_text(request: PromptRequest):
     
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"response": response}
-

requirements.txt

@@ -3,4 +3,6 @@ transformers
 torch
 uvicorn
 python-dotenv
-accelerate>=0.26.0
+optimum[onnxruntime]  # For CPU optimizations with ONNX Runtime
+accelerate            # For managing multi-device setup (CPU/GPU)
+gunicorn              # For running multiple workers