Use new definitions of Model / SKU

llama_toolchain/inference/api/endpoints.py

@@ -13,7 +13,7 @@ from pyopenapi import webmethod
 
 @json_schema_type
 class CompletionRequest(BaseModel):
-    model: PretrainedModel
+    model: str
     content: InterleavedTextAttachment
     sampling_params: Optional[SamplingParams] = SamplingParams()
 
@@ -39,7 +39,7 @@ class CompletionResponseStreamChunk(BaseModel):
 
 @json_schema_type
 class BatchCompletionRequest(BaseModel):
-    model: PretrainedModel
+    model: str
     content_batch: List[InterleavedTextAttachment]
     sampling_params: Optional[SamplingParams] = SamplingParams()
     logprobs: Optional[LogProbConfig] = None
@@ -53,7 +53,7 @@ class BatchCompletionResponse(BaseModel):
 
 @json_schema_type
 class ChatCompletionRequest(BaseModel):
-    model: InstructModel
+    model: str
     messages: List[Message]
     sampling_params: Optional[SamplingParams] = SamplingParams()
 
@@ -80,7 +80,7 @@ class ChatCompletionResponse(BaseModel):
 
 @json_schema_type
 class BatchChatCompletionRequest(BaseModel):
-    model: InstructModel
+    model: str
     messages_batch: List[List[Message]]
     sampling_params: Optional[SamplingParams] = SamplingParams()
 

llama_toolchain/inference/client.py

@@ -18,7 +18,6 @@ from .api import (
     ChatCompletionResponseStreamChunk,
     CompletionRequest,
     Inference,
-    InstructModel,
     UserMessage,
 )
 from .event_logger import EventLogger
@@ -67,7 +66,7 @@ async def run_main(host: str, port: int, stream: bool):
     cprint(f"User>{message.content}", "green")
     iterator = client.chat_completion(
         ChatCompletionRequest(
-            model=InstructModel.llama3_8b_chat,
+            model="Meta-Llama-3.1-8B-Instruct",
             messages=[message],
             stream=stream,
         )