mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-18 12:10:34 +08:00
Yannic KilcherandGitHub
1709dc0324
* very primitive implementation of inference * re-worked with security in mind * removed polling from clients * switched workers to websockets * implemented back and forth chats
22 lines
499 B
Python
22 lines
499 B
Python
import random
|
|
|
|
import pydantic
|
|
|
|
from . import protocol
|
|
|
|
|
|
class WorkerConfig(pydantic.BaseModel):
|
|
model_name: str = "distilgpt2"
|
|
|
|
|
|
class WorkRequest(pydantic.BaseModel):
|
|
conversation: protocol.Conversation = pydantic.Field(..., repr=False)
|
|
model_name: str = "distilgpt2"
|
|
max_new_tokens: int = 100
|
|
seed: int = pydantic.Field(default_factory=lambda: random.randint(0, 2**32 - 1))
|
|
|
|
|
|
class WorkResponsePacket(pydantic.BaseModel):
|
|
token: str | None = None
|
|
is_end: bool = False
|