Update main.py
Browse files
main.py
CHANGED
@@ -1,10 +1,8 @@
|
|
1 |
from typing import Union
|
2 |
-
|
3 |
from fastapi import FastAPI
|
4 |
-
|
5 |
from huggingface_hub import InferenceClient
|
6 |
-
|
7 |
import os
|
|
|
8 |
|
9 |
app = FastAPI()
|
10 |
|
@@ -13,11 +11,11 @@ async def root():
|
|
13 |
return {"message": "Hello World"}
|
14 |
|
15 |
@app.post("/chat")
|
16 |
-
async def chat(
|
17 |
try:
|
18 |
client = InferenceClient(model="gpt2", token=os.environ.get(api_key))
|
19 |
-
input_text =
|
20 |
-
response = client.text_generation(input_text
|
21 |
print(response)
|
22 |
return {"message": response}
|
23 |
except Exception as e:
|
|
|
1 |
from typing import Union
|
|
|
2 |
from fastapi import FastAPI
|
|
|
3 |
from huggingface_hub import InferenceClient
|
|
|
4 |
import os
|
5 |
+
import json
|
6 |
|
7 |
app = FastAPI()
|
8 |
|
|
|
11 |
return {"message": "Hello World"}
|
12 |
|
13 |
@app.post("/chat")
|
14 |
+
async def chat(message: dict):
|
15 |
try:
|
16 |
client = InferenceClient(model="gpt2", token=os.environ.get(api_key))
|
17 |
+
input_text = message
|
18 |
+
response = client.text_generation(input_text)
|
19 |
print(response)
|
20 |
return {"message": response}
|
21 |
except Exception as e:
|