beet_ml 0.0.8

Machine Learning actions built upon beet_flow
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

use serde::Deserialize;
use serde::Serialize;


#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BertConfig {
	pub model: BertModelConfig,
	pub normalize_embeddings: bool,
	pub approximate_gelu: bool,
}

impl Default for BertConfig {
	fn default() -> Self {
		Self{
			model: BertModelConfig {
					base_url: "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/refs%2Fpr%2F21/".into(),
					search_prefix: "".into(),
					document_prefix: "".into(),
					model_id: "sentence-transformers/all-MiniLM-L6-v2".into(),
					revision: "refs/pr/21".into(),
		},
			normalize_embeddings: true,
			approximate_gelu: false,
		}
	}
}

impl BertConfig {
	pub fn new(model: BertModelConfig) -> Self {
		Self {
			model,
			normalize_embeddings: true,
			approximate_gelu: false,
		}
	}
}


/// Config containing both native and wasm urls for a model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BertModelConfig {
	pub base_url: String,
	pub search_prefix: String,
	pub document_prefix: String,
	/// used with hf-hub
	pub model_id: String,
	pub revision: String,
}


impl BertModelConfig {
	pub fn model_url(&self) -> String {
		self.base_url.clone() + "model.safetensors"
	}

	pub fn config_url(&self) -> String { self.base_url.clone() + "config.json" }

	pub fn tokenizer_url(&self) -> String {
		// rejected 403 in github actions, use s3 instead
		// self.base_url.clone() + "tokenizer.json"
		"https://bevyhub-public.s3.us-west-2.amazonaws.com/assets/ml/tokenizer.json".into()
	}
}