David commited on
Commit
bed79e4
·
1 Parent(s): b9af006

refactor: update InferenceClient configuration and specify llama model in chat completion request

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -6,8 +6,8 @@ import os
6
  app = FastAPI()
7
 
8
  client = InferenceClient(
9
- model="nvidia/Llama-3.1-Nemotron-Nano-8B-v1",
10
- token=os.environ.get("HF_TOKEN")
11
  )
12
 
13
  class ChatRequest(BaseModel):
@@ -20,6 +20,7 @@ def greet_json():
20
  @app.post("/chat")
21
  def chat(request: ChatRequest):
22
  response = client.chat.completions.create(
 
23
  messages=[{"role": "user", "content": request.message}],
24
  max_tokens=1024,
25
  )
 
6
  app = FastAPI()
7
 
8
  client = InferenceClient(
9
+ provider="hf-inference",
10
+ api_key=os.environ.get("HF_TOKEN")
11
  )
12
 
13
  class ChatRequest(BaseModel):
 
20
  @app.post("/chat")
21
  def chat(request: ChatRequest):
22
  response = client.chat.completions.create(
23
+ model="meta-llama/Llama-3.1-8B-Instruct",
24
  messages=[{"role": "user", "content": request.message}],
25
  max_tokens=1024,
26
  )