Update main.py

main.py

@@ -1,5 +1,5 @@
 import os
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends
 from pydantic import BaseModel
 from ctransformers import AutoModelForCausalLM
 
@@ -27,7 +27,7 @@ def setup_endpoints(app):
 
     models = {}
 
-    # Load each model and create an endpoint
+    # Load each model
     for model_name in model_dirs:
         model_path = os.path.join(model_base_path, model_name)
         try:
@@ -37,8 +37,16 @@ def setup_endpoints(app):
             print(f"Failed to load model {model_name}: {e}")
             continue
 
+    # Function to get model dependency
+    def get_model(model_name: str):
+        if model_name not in models:
+            raise HTTPException(status_code=404, detail="Model not found")
+        return models[model_name]
+
+    # Create an endpoint for each model
+    for model_name in model_dirs:
         @app.post(f"/{model_name}")
-        async def generate_response(item: Validation, model=model):
+        async def generate_response(item: Validation, model=Depends(lambda: get_model(model_name))):
             try:
                 response = model(item.inputs,
                                  temperature=item.temperature,