172 lines
3.8 KiB
Python
172 lines
3.8 KiB
Python
from pydantic import BaseModel
|
|
from typing import List, Literal
|
|
|
|
|
|
class MultipleChoiceSchema(BaseModel):
|
|
answer: Literal["A", "B", "C", "D"]
|
|
|
|
|
|
class ListOfNumbersSchema(BaseModel):
|
|
answer: List[int]
|
|
|
|
|
|
class ListofStringsSchema(BaseModel):
|
|
answer: List[str]
|
|
|
|
|
|
class NumberSchema(BaseModel):
|
|
answer: int
|
|
|
|
|
|
class StringSchema(BaseModel):
|
|
answer: str
|
|
|
|
|
|
# Winogrande Models #############################
|
|
|
|
|
|
class BinaryChoiceSchema(BaseModel):
|
|
answer: Literal["A", "B"]
|
|
|
|
|
|
# BBQ Models #############################
|
|
|
|
|
|
class TrinaryChoiceSchema(BaseModel):
|
|
answer: Literal["A", "B", "C"]
|
|
|
|
|
|
# MathQA Models #############################
|
|
|
|
|
|
class MultipleChoiceSchemaLower(BaseModel):
|
|
answer: Literal["a", "b", "c", "d"]
|
|
|
|
|
|
# DROP Models #############################
|
|
|
|
|
|
class DROPStringSchema(BaseModel):
|
|
answer: str
|
|
|
|
|
|
class DROPNumberSchema(BaseModel):
|
|
answer: int
|
|
|
|
|
|
class DROPDateSchema(BaseModel):
|
|
answer: str
|
|
|
|
|
|
# BBH Models #############################
|
|
|
|
|
|
class AffirmationSchema(BaseModel):
|
|
answer: Literal["No", "Yes"]
|
|
|
|
|
|
class AffirmationLowerSchema(BaseModel):
|
|
answer: Literal["no", "yes"]
|
|
|
|
|
|
class BooleanSchema(BaseModel):
|
|
answer: Literal["True", "False"]
|
|
|
|
|
|
class ValidSchema(BaseModel):
|
|
answer: Literal["valid", "invalid"]
|
|
|
|
|
|
class BBHMultipleChoice2Schema(BaseModel):
|
|
answer: Literal["(A)", "(B)"]
|
|
|
|
|
|
class BBHMultipleChoice3Schema(BaseModel):
|
|
answer: Literal["(A)", "(B)", "(C)"]
|
|
|
|
|
|
class BBHMultipleChoice4Schema(BaseModel):
|
|
answer: Literal["(A)", "(B)", "(C)", "(D)"]
|
|
|
|
|
|
class BBHMultipleChoice5Schema(BaseModel):
|
|
answer: Literal["(A)", "(B)", "(C)", "(D)", "(E)"]
|
|
|
|
|
|
class BBHMultipleChoice6Schema(BaseModel):
|
|
answer: Literal["(A)", "(B)", "(C)", "(D)", "(E)", "(F)"]
|
|
|
|
|
|
class BBHMultipleChoice7Schema(BaseModel):
|
|
answer: Literal["(A)", "(B)", "(C)", "(D)", "(E)", "(F)", "(G)"]
|
|
|
|
|
|
class BBHMultipleChoice11Schema(BaseModel):
|
|
answer: Literal[
|
|
"(A)",
|
|
"(B)",
|
|
"(C)",
|
|
"(D)",
|
|
"(E)",
|
|
"(F)",
|
|
"(G)",
|
|
"(H)",
|
|
"(I)",
|
|
"(J)",
|
|
"(K)",
|
|
]
|
|
|
|
|
|
class BBHMultipleChoice18Schema(BaseModel):
|
|
answer: Literal[
|
|
"(A)",
|
|
"(B)",
|
|
"(C)",
|
|
"(D)",
|
|
"(E)",
|
|
"(F)",
|
|
"(G)",
|
|
"(H)",
|
|
"(I)",
|
|
"(J)",
|
|
"(K)",
|
|
"(L)",
|
|
"(M)",
|
|
"(N)",
|
|
"(O)",
|
|
"(P)",
|
|
"(Q)",
|
|
"(R)",
|
|
]
|
|
|
|
|
|
bbh_models_dict = {
|
|
"boolean_expressions": BooleanSchema,
|
|
"causal_judgement": AffirmationSchema,
|
|
"date_understanding": BBHMultipleChoice6Schema,
|
|
"disambiguation_qa": BBHMultipleChoice3Schema,
|
|
"dyck_languages": StringSchema,
|
|
"formal_fallacies": ValidSchema,
|
|
"geometric_shapes": BBHMultipleChoice11Schema,
|
|
"hyperbaton": BBHMultipleChoice2Schema,
|
|
"logical_deduction_three_objects": BBHMultipleChoice3Schema,
|
|
"logical_deduction_five_objects": BBHMultipleChoice5Schema,
|
|
"logical_deduction_seven_objects": BBHMultipleChoice7Schema,
|
|
"movie_recommendation": BBHMultipleChoice5Schema,
|
|
"multistep_arithmetic_two": NumberSchema,
|
|
"navigate": AffirmationSchema,
|
|
"object_counting": NumberSchema,
|
|
"penguins_in_a_table": BBHMultipleChoice5Schema,
|
|
"reasoning_about_colored_objects": BBHMultipleChoice18Schema,
|
|
"ruin_names": BBHMultipleChoice4Schema,
|
|
"salient_translation_error_detection": BBHMultipleChoice6Schema,
|
|
"snarks": BBHMultipleChoice2Schema,
|
|
"sports_understanding": AffirmationLowerSchema,
|
|
"temporal_sequences": BBHMultipleChoice4Schema,
|
|
"tracking_shuffled_objects_three_objects": BBHMultipleChoice3Schema,
|
|
"tracking_shuffled_objects_five_objects": BBHMultipleChoice5Schema,
|
|
"tracking_shuffled_objects_seven_objects": BBHMultipleChoice7Schema,
|
|
"web_of_lies": AffirmationSchema,
|
|
"word_sorting": StringSchema,
|
|
}
|