Skip to content

Commit

Permalink
Update openai_api.py
Browse files Browse the repository at this point in the history
  • Loading branch information
hiyouga authored Jun 26, 2023
1 parent 9edf43f commit 892691c
Showing 1 changed file with 23 additions and 1 deletion.
24 changes: 23 additions & 1 deletion openai_api.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# coding=utf-8
# Implements API for ChatGLM2-6B in OpenAI's format. (https://platform.openai.com/docs/api-reference/chat)
# Usage: python openai_api.py
# Usage: python api_demo.py
# Visit http://localhost:8000/docs for documents.


Expand All @@ -26,6 +26,21 @@ async def lifespan(app: FastAPI): # collects GPU memory
app = FastAPI(lifespan=lifespan)


class ModelCard(BaseModel):
id: str
object: str = "model"
created: int = Field(default_factory=lambda: int(time.time()))
owned_by: str = "owner"
root: Optional[str] = None
parent: Optional[str] = None
permission: Optional[list] = None


class ModelList(BaseModel):
object: str = "list"
data: List[ModelCard] = []


class ChatMessage(BaseModel):
role: Literal["user", "assistant", "system"]
content: str
Expand Down Expand Up @@ -64,6 +79,13 @@ class ChatCompletionResponse(BaseModel):
created: Optional[int] = Field(default_factory=lambda: int(time.time()))


@app.get("/v1/models", response_model=ModelList)
async def list_models():
global model_args
model_card = ModelCard(id="gpt-3.5-turbo")
return ModelList(data=[model_card])


@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
async def create_chat_completion(request: ChatCompletionRequest):
global model, tokenizer
Expand Down

0 comments on commit 892691c

Please sign in to comment.