Update app.py

app.py

@@ -123,8 +123,19 @@ def build_api_messages(session_id: int, system_message: str):
     msgs.extend(get_messages(session_id))
     return msgs
 
-def get_client(model_choice: str):
-    return InferenceClient(MODELS.get(model_choice, list(MODELS.values())[0]), token=HF_TOKEN)
+# def get_client(model_choice: str):
+#     return InferenceClient(MODELS.get(model_choice, list(MODELS.values())[0]), token=HF_TOKEN)
+def get_client(model_choice):
+    # Normalize model_choice -> must be a string
+    if isinstance(model_choice, list):
+        model_choice = model_choice[0] if model_choice else None
+
+    if not model_choice:
+        model_choice = list(MODELS.keys())[0]
+
+    model_id = MODELS.get(model_choice, list(MODELS.values())[0])
+    return InferenceClient(model_id, token=HF_TOKEN)
+
 
 def load_dataset_by_name(name: str):
     if name == "The Stack": return load_dataset("bigcode/the-stack", split="train")