<commit_before>package word2vec

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
)

type cosQuery struct {
	A Expr `json:"a,omitempty"`
	B Expr `json:"b,omitempty"`
}

type cosResponse struct {
	Value float32 `json:"value"`
}

func (q cosQuery) Eval(m *Model) (*cosResponse, error) {
	v, err := m.Cos(q.A, q.B)
	if err != nil {
		return nil, err
	}

	return &cosResponse{
		Value: v,
	}, nil
}

type cosesQuery struct {
	Queries []cosQuery `json:"queries"`
}

type cosesResponse struct {
	Values []cosResponse `json:"values"`
}

func (qs cosesQuery) Eval(m *Model) (*cosesResponse, error) {
	values := make([]cosResponse, len(qs.Queries))
	for i, q := range qs.Queries {
		r, err := q.Eval(m)
		if err != nil {
			return nil, err
		}
		values[i] = *r
	}
	return &cosesResponse{
		Values: values,
	}, nil
}

type cosNQuery struct {
	Expr Expr `json:"expr"`
	N    int  `json:"n"`
}

type cosNResponse struct {
	Matches []Match `json:"matches"`
}

func (q cosNQuery) Eval(m *Model) (*cosNResponse, error) {
	r, err := m.CosN(q.Expr, q.N)
	if err != nil {
		return nil, err
	}

	return &cosNResponse{
		Matches: r,
	}, nil
}

// Server is a type which implements http.Handler and exports endpoints
// for performing similarity queries on a word2vec model.
type Server struct {
	*Model
	*http.ServeMux
}

// NewServer creates a new word2vec Server which exports endpoints for performing
// similarity queries on a word2vec Model.
func NewServer(m *Model) http.Handler {
	ms := &Server{
		Model: m,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/cos-n", ms.handleCosNQuery)
	mux.HandleFunc("/cos", ms.handleCosQuery)
	mux.HandleFunc("/coses", ms.handleCosesQuery)

	ms.ServeMux = mux
	return ms
}

func handleError(w http.ResponseWriter, r *http.Request, status int, msg string) {
	log.Printf(msg)
	w.WriteHeader(status)
	w.Write([]byte(msg))
	return
}

func (m *Server) handleCosQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q cosQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *Server) handleCosesQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q cosesQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *Server) handleCosNQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q cosNQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}
<commit_msg>Un-export server type, only need NewServer method to be exported.<commit_after>package word2vec

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
)

type cosQuery struct {
	A Expr `json:"a,omitempty"`
	B Expr `json:"b,omitempty"`
}

type cosResponse struct {
	Value float32 `json:"value"`
}

func (q cosQuery) Eval(m *Model) (*cosResponse, error) {
	v, err := m.Cos(q.A, q.B)
	if err != nil {
		return nil, err
	}

	return &cosResponse{
		Value: v,
	}, nil
}

type cosesQuery struct {
	Queries []cosQuery `json:"queries"`
}

type cosesResponse struct {
	Values []cosResponse `json:"values"`
}

func (qs cosesQuery) Eval(m *Model) (*cosesResponse, error) {
	values := make([]cosResponse, len(qs.Queries))
	for i, q := range qs.Queries {
		r, err := q.Eval(m)
		if err != nil {
			return nil, err
		}
		values[i] = *r
	}
	return &cosesResponse{
		Values: values,
	}, nil
}

type cosNQuery struct {
	Expr Expr `json:"expr"`
	N    int  `json:"n"`
}

type cosNResponse struct {
	Matches []Match `json:"matches"`
}

func (q cosNQuery) Eval(m *Model) (*cosNResponse, error) {
	r, err := m.CosN(q.Expr, q.N)
	if err != nil {
		return nil, err
	}

	return &cosNResponse{
		Matches: r,
	}, nil
}

// server is a type which implements http.Handler and exports endpoints
// for performing similarity queries on a word2vec model.
type server struct {
	*Model
	*http.ServeMux
}

// NewServer creates a new word2vec server which exports endpoints for performing
// similarity queries on a word2vec Model.
func NewServer(m *Model) http.Handler {
	ms := &server{
		Model: m,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/cos-n", ms.handleCosNQuery)
	mux.HandleFunc("/cos", ms.handleCosQuery)
	mux.HandleFunc("/coses", ms.handleCosesQuery)

	ms.ServeMux = mux
	return ms
}

func handleError(w http.ResponseWriter, r *http.Request, status int, msg string) {
	log.Printf(msg)
	w.WriteHeader(status)
	w.Write([]byte(msg))
	return
}

func (m *server) handleCosQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q cosQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *server) handleCosesQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q cosesQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *server) handleCosNQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q cosNQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

{'fraction_non_alphanumeric': 0.15167064439140812, 'fraction_numerical': 0.001431980906921241, 'mean_word_length': 3.844508670520231, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3410371', 'n_tokens_mistral': 3259, 'n_tokens_neox': 3115, 'n_words': 875}