Feat: bench

main.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"grailbench/config"
+	"grailbench/models"
 )
 
 var (
@@ -20,26 +21,13 @@ var (
 	currentModel = ""
 )
 
-type Question struct {
-	ID       string `json:"id"`
-	Topic    string `json:"topic"`
-	Question string `json:"question"`
-}
-
-type Answer struct {
-	Q      Question
-	Answer string `json:"answer"`
-	Model  string `json:"model"`
-	// resp time?
-}
-
-func loadQuestions(fp string) ([]Question, error) {
+func loadQuestions(fp string) ([]models.Question, error) {
 	data, err := os.ReadFile(fp)
 	if err != nil {
 		logger.Error("failed to read file", "error", err, "fp", fp)
 		return nil, err
 	}
-	resp := []Question{}
+	resp := []models.Question{}
 	if err := json.Unmarshal(data, &resp); err != nil {
 		logger.Error("failed to unmarshal file", "error", err, "fp", fp)
 		return nil, err
@@ -77,8 +65,8 @@ func main() {
 	}
 }
 
-func runBench(questions []Question) ([]Answer, error) {
-	answers := []Answer{}
+func runBench(questions []models.Question) ([]models.Answer, error) {
+	answers := []models.Answer{}
 	for _, q := range questions {
 		resp, err := callLLM(buildPrompt(q.Question))
 		if err != nil {
@@ -90,7 +78,7 @@ func runBench(questions []Question) ([]Answer, error) {
 		if err != nil {
 			panic(err)
 		}
-		a := Answer{Q: q, Answer: respText, Model: currentModel}
+		a := models.Answer{Q: q, Answer: respText, Model: currentModel}
 		answers = append(answers, a)
 	}
 	return answers, nil