mirror of https://github.com/HazyResearch/manifest
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
"""Response."""
|
|
|
|
import time
|
|
import uuid
|
|
from typing import Any, Dict, List
|
|
|
|
|
|
class ModelResponse:
|
|
"""ModelResponse."""
|
|
|
|
def __init__(self, results: List[Dict[str, Any]], response_type: str) -> None:
|
|
"""Initialize response."""
|
|
self.results = results
|
|
self.response_type = response_type
|
|
if self.response_type not in {
|
|
"text_completion",
|
|
"prompt_logit_score",
|
|
"image_generation",
|
|
}:
|
|
raise ValueError(
|
|
f"Invalid response type: {self.response_type}. "
|
|
"Must be one of: text_completion, prompt_logit_score, image_generation."
|
|
)
|
|
self.response_id = str(uuid.uuid4())
|
|
self.created = int(time.time())
|
|
|
|
def __dict__(self) -> Dict[str, Any]: # type: ignore
|
|
"""Return dictionary representation of response."""
|
|
key = "text" if self.response_type != "image_generation" else "array"
|
|
return {
|
|
"id": self.response_id,
|
|
"object": self.response_type,
|
|
"created": self.created,
|
|
"model": "flask_model",
|
|
"choices": [
|
|
{
|
|
key: result[key],
|
|
"logprob": result["logprob"],
|
|
}
|
|
if key == "text"
|
|
else {
|
|
key: result[key].tolist(),
|
|
"logprob": result["logprob"],
|
|
}
|
|
for result in self.results
|
|
],
|
|
}
|