mgbam commited on
Commit
1edddce
·
verified ·
1 Parent(s): 01ecbcf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -3
app.py CHANGED
@@ -215,24 +215,28 @@ DEMO_LIST = [
215
  }
216
  ]
217
 
218
- # HF Inference Client
219
  HF_TOKEN = os.getenv('HF_TOKEN')
220
  GROQ_API_KEY = os.getenv('GROQ_API_KEY')
 
221
 
222
  def get_inference_client(model_id):
223
- """Return an InferenceClient configured for Hugging Face or Groq."""
224
  if model_id == "moonshotai/Kimi-K2-Instruct":
225
  return InferenceClient(
226
  base_url="https://api.groq.com/openai/v1",
227
  api_key=GROQ_API_KEY
228
  )
 
 
 
 
 
229
  else:
230
  return InferenceClient(
231
  model=model_id,
232
  api_key=HF_TOKEN
233
  )
234
 
235
-
236
  # Type definitions
237
  History = List[Tuple[str, str]]
238
  Messages = List[Dict[str, str]]
 
215
  }
216
  ]
217
 
 
218
  HF_TOKEN = os.getenv('HF_TOKEN')
219
  GROQ_API_KEY = os.getenv('GROQ_API_KEY')
220
+ FIREWORKS_API_KEY = os.getenv('FIREWORKS_API_KEY')
221
 
222
  def get_inference_client(model_id):
223
+ """Return an InferenceClient configured for Hugging Face, Groq, or Fireworks AI."""
224
  if model_id == "moonshotai/Kimi-K2-Instruct":
225
  return InferenceClient(
226
  base_url="https://api.groq.com/openai/v1",
227
  api_key=GROQ_API_KEY
228
  )
229
+ elif model_id.startswith("fireworks/"):
230
+ return InferenceClient(
231
+ base_url="https://api.fireworks.ai/inference/v1",
232
+ api_key=FIREWORKS_API_KEY
233
+ )
234
  else:
235
  return InferenceClient(
236
  model=model_id,
237
  api_key=HF_TOKEN
238
  )
239
 
 
240
  # Type definitions
241
  History = List[Tuple[str, str]]
242
  Messages = List[Dict[str, str]]