yusufs commited on
Commit
b41be20
·
1 Parent(s): 35decf8

feat(response): should dict only

Browse files
Files changed (1) hide show
  1. main.py +9 -2
main.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import random
2
  import torch
3
  from typing import Any
@@ -121,7 +122,7 @@ class GenerationResponse(BaseModel):
121
 
122
 
123
  @app.post("/generate-llama3-2")
124
- def generate_text(request: GenerationRequest) -> list[RequestOutput] | dict[str, str]:
125
  try:
126
  sampling_params: SamplingParams = SamplingParams(
127
  temperature=request.temperature,
@@ -130,11 +131,17 @@ def generate_text(request: GenerationRequest) -> list[RequestOutput] | dict[str,
130
  )
131
 
132
  # Generate text
133
- return engine_llama_3_2.generate(
134
  prompts=request.prompt,
135
  sampling_params=sampling_params
136
  )
137
 
 
 
 
 
 
 
138
  except Exception as e:
139
  return {
140
  "error": str(e)
 
1
+ import json
2
  import random
3
  import torch
4
  from typing import Any
 
122
 
123
 
124
  @app.post("/generate-llama3-2")
125
+ def generate_text(request: GenerationRequest) -> dict[str, Any]:
126
  try:
127
  sampling_params: SamplingParams = SamplingParams(
128
  temperature=request.temperature,
 
131
  )
132
 
133
  # Generate text
134
+ response: list[RequestOutput] = engine_llama_3_2.generate(
135
  prompts=request.prompt,
136
  sampling_params=sampling_params
137
  )
138
 
139
+ output: str = json.dumps(obj=response)
140
+
141
+ return {
142
+ "output": json.loads(s=output),
143
+ }
144
+
145
  except Exception as e:
146
  return {
147
  "error": str(e)