Update main.py

main.py

@@ -25,7 +25,7 @@ os.makedirs(model_dir, exist_ok=True)
 recog_path_local = api.hf_hub_download(repo_id="KennethTM/jms-model", filename="recog_model.onnx", local_dir=model_dir, repo_type="model")
 corner_path_local = api.hf_hub_download(repo_id="KennethTM/jms-model", filename="corner_model.onnx", local_dir=model_dir, repo_type="model")
 card_data_path = api.hf_hub_download(repo_id="KennethTM/jms-model", filename="card_data_minimal.parquet", local_dir=model_dir, repo_type="model")
-card_embeddings_path = api.hf_hub_download(repo_id="KennethTM/jms-model", filename="card_embeddings.npz", local_dir=model_dir, repo_type="model")
+card_embeddings_path = api.hf_hub_download(repo_id="KennethTM/jms-model", filename="card_embeddings_float16.npz", local_dir=model_dir, repo_type="model")
 task_config_path = api.hf_hub_download(repo_id="KennethTM/jms-model", filename="task_config.json", local_dir=model_dir, repo_type="model")
 
 # Initialize FastAPI app
@@ -41,10 +41,10 @@ recog_session = ort.InferenceSession(recog_path_local)
 
 # Load reference embeddings and card data
 df = pd.read_parquet(card_data_path)
-ref_embeddings = np.load(card_embeddings_path)['embeddings']
+ref_embeddings = np.load(card_embeddings_path)['embeddings'].astype(np.float32)
 
 # Pre-compute card info as list of dicts for faster access (avoid DataFrame iloc overhead)
-card_metadata = df[['card_id', 'name', 'uri', 'card_url', 'image_url', 'lang', 'rarity', 'set_name', 'set']].to_dict('records')
+card_metadata = df[['name', 'card_url', 'image_url', 'rarity', 'set_name', 'set']].to_dict('records')
 del df  # Free DataFrame memory after extracting needed data
 
 with open(task_config_path) as f:
@@ -115,12 +115,9 @@ def preprocess_onnx(image: np.ndarray, task: str) -> np.ndarray:
     return image
 
 class Card(BaseModel):
-    id: str
     name: str
-    uri: str
     scryfall_uri: str
     image_url: str
-    lang: str
     rarity: str
     set_name: str
     set: str
@@ -191,12 +188,9 @@ async def predict(file: UploadFile = File(...)) -> Card:
     prediction_time_ms = int((t1 - t0) * 1000)
 
     return Card(
-        id=card_info['card_id'],
         name=card_info['name'],
-        uri=card_info['uri'],
         scryfall_uri=card_info['card_url'],
         image_url=card_info['image_url'],
-        lang=card_info['lang'],
         rarity=card_info['rarity'],
         set_name=card_info['set_name'],
         set=card_info['set'],