André Fernandes commited on
Commit
6c6abf4
·
1 Parent(s): c876c9f

added methods to get specific hosted models from API

Browse files
Files changed (1) hide show
  1. app.py +38 -2
app.py CHANGED
@@ -1,7 +1,43 @@
1
- from fastapi import FastAPI
 
 
 
 
 
 
 
 
 
 
2
 
3
  app = FastAPI()
4
 
 
 
 
 
 
 
 
 
 
5
  @app.get("/")
6
  def greet_json():
7
- return {"Hello": "World!"}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Optional
2
+
3
+ from fastapi import FastAPI, HTTPException
4
+ from pydantic import BaseModel
5
+
6
+
7
+ class Model(BaseModel):
8
+ id: int
9
+ name: str
10
+ param_count: Optional[int] = None
11
+
12
 
13
  app = FastAPI()
14
 
15
+ models = [Model(id=0, name="CNN"), Model(id=1, name="Transformer")]
16
+ id_2_hosted_models = {
17
+ model.id : model for model in models
18
+ }
19
+ model_names_2_id = {
20
+ model.name.lower() : model.id for model in models
21
+ }
22
+
23
+
24
  @app.get("/")
25
  def greet_json():
26
+ return {"Hello World": "Welcome to my ML Repository API!"}
27
+
28
+ @app.get("/hosted/id/{model_id}")
29
+ def get_by_id(model_id: int):
30
+
31
+ if model_id not in id_2_hosted_models:
32
+ raise HTTPException(status_code=404, detail=f"Model with id={model_id} not found")
33
+
34
+ return id_2_hosted_models[model_id]
35
+
36
+ @app.get("/hosted/name/{model_name}")
37
+ def get_by_name(model_name: str):
38
+ model_name = model_name.lower()
39
+
40
+ if model_name not in model_names_2_id:
41
+ raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found")
42
+
43
+ return id_2_hosted_models[model_names_2_id[model_name]]