@sjcrh/proteinpaint-rust 2.167.0 → 2.170.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/src/test_ai.rs CHANGED
@@ -5,6 +5,7 @@ fn main() {}
5
5
 
6
6
  #[cfg(test)]
7
7
  mod tests {
8
+ use crate::aichatbot::{AiJsonFormat, AnswerFormat, SummaryType, llm_backend, run_pipeline};
8
9
  use serde_json;
9
10
  use std::fs::{self};
10
11
  use std::path::Path;
@@ -20,6 +21,7 @@ mod tests {
20
21
  ollama_comp_model_name: String,
21
22
  ollama_embedding_model_name: String,
22
23
  genomes: Vec<Genomes>,
24
+ aiRoute: String,
23
25
  }
24
26
 
25
27
  #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
@@ -49,7 +51,7 @@ mod tests {
49
51
 
50
52
  // Parse the JSON data
51
53
  let serverconfig: ServerConfig = serde_json::from_str(&data).expect("JSON not in serverconfig.json format");
52
-
54
+ let airoute = String::from("../../") + &serverconfig.aiRoute;
53
55
  for genome in &serverconfig.genomes {
54
56
  for dataset in &genome.datasets {
55
57
  match &dataset.aifiles {
@@ -61,13 +63,13 @@ mod tests {
61
63
  // Read the file
62
64
  let ai_data = fs::read_to_string(ai_json_file).unwrap();
63
65
  // Parse the JSON data
64
- let ai_json: super::super::AiJsonFormat =
66
+ let ai_json: AiJsonFormat =
65
67
  serde_json::from_str(&ai_data).expect("AI JSON file does not have the correct format");
66
68
  //println!("ai_json:{:?}", ai_json);
67
69
  let genedb = String::from(&serverconfig.tpmasterdir) + &"/" + &ai_json.genedb;
68
70
  let dataset_db = String::from(&serverconfig.tpmasterdir) + &"/" + &ai_json.db;
69
71
  let llm_backend_name = &serverconfig.llm_backend;
70
- let llm_backend_type: super::super::llm_backend;
72
+ let llm_backend_type: llm_backend;
71
73
 
72
74
  if llm_backend_name != "ollama" && llm_backend_name != "SJ" {
73
75
  panic!(
@@ -77,7 +79,7 @@ mod tests {
77
79
  let ollama_host = &serverconfig.ollama_apilink;
78
80
  let ollama_embedding_model_name = &serverconfig.ollama_embedding_model_name;
79
81
  let ollama_comp_model_name = &serverconfig.ollama_comp_model_name;
80
- llm_backend_type = super::super::llm_backend::Ollama();
82
+ llm_backend_type = llm_backend::Ollama();
81
83
  let ollama_client = super::super::ollama::Client::builder()
82
84
  .base_url(ollama_host)
83
85
  .build()
@@ -85,40 +87,45 @@ mod tests {
85
87
  let embedding_model = ollama_client.embedding_model(ollama_embedding_model_name);
86
88
  let comp_model = ollama_client.completion_model(ollama_comp_model_name);
87
89
  for chart in ai_json.charts.clone() {
88
- match chart {
89
- super::super::Charts::Summary(testdata) => {
90
- for ques_ans in testdata.TestData {
91
- let user_input = ques_ans.question;
92
- let llm_output = super::super::run_pipeline(
93
- &user_input,
94
- comp_model.clone(),
95
- embedding_model.clone(),
96
- llm_backend_type.clone(),
97
- temperature,
98
- max_new_tokens,
99
- top_p,
100
- &dataset_db,
101
- &genedb,
102
- &ai_json,
103
- testing,
104
- )
105
- .await;
106
- let mut llm_json_value: super::super::SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
107
- let mut expected_json_value: super::super::SummaryType = serde_json::from_str(&ques_ans.answer).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
108
- assert_eq!(
109
- llm_json_value.sort_summarytype_struct(),
110
- expected_json_value.sort_summarytype_struct()
111
- );
90
+ if chart.r#type == "Summary" {
91
+ for ques_ans in chart.TestData {
92
+ let user_input = ques_ans.question;
93
+ let llm_output = run_pipeline(
94
+ &user_input,
95
+ comp_model.clone(),
96
+ embedding_model.clone(),
97
+ llm_backend_type.clone(),
98
+ temperature,
99
+ max_new_tokens,
100
+ top_p,
101
+ &dataset_db,
102
+ &genedb,
103
+ &ai_json,
104
+ &airoute,
105
+ testing,
106
+ )
107
+ .await;
108
+ let llm_json_value: SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
109
+ match ques_ans.answer {
110
+ AnswerFormat::summary_type(sum) => {
111
+ //println!("expected answer:{:?}", &sum);
112
+ assert_eq!(
113
+ llm_json_value.sort_summarytype_struct(),
114
+ sum.sort_summarytype_struct()
115
+ );
116
+ }
117
+ AnswerFormat::DE_type(_) => {
118
+ panic!("DE type not valid for summary")
119
+ }
112
120
  }
113
121
  }
114
- super::super::Charts::DE(_testdata) => {} // To do
115
122
  }
116
123
  }
117
124
  } else if *llm_backend_name == "SJ".to_string() {
118
125
  let sjprovider_host = &serverconfig.sj_apilink;
119
126
  let sj_embedding_model_name = &serverconfig.sj_embedding_model_name;
120
127
  let sj_comp_model_name = &serverconfig.sj_comp_model_name;
121
- llm_backend_type = super::super::llm_backend::Sj();
128
+ llm_backend_type = llm_backend::Sj();
122
129
  let sj_client = super::super::sjprovider::Client::builder()
123
130
  .base_url(sjprovider_host)
124
131
  .build()
@@ -127,37 +134,53 @@ mod tests {
127
134
  let comp_model = sj_client.completion_model(sj_comp_model_name);
128
135
 
129
136
  for chart in ai_json.charts.clone() {
130
- match chart {
131
- super::super::Charts::Summary(testdata) => {
132
- for ques_ans in testdata.TestData {
133
- let user_input = ques_ans.question;
134
- if user_input.len() > 0 {
135
- let llm_output = super::super::run_pipeline(
136
- &user_input,
137
- comp_model.clone(),
138
- embedding_model.clone(),
139
- llm_backend_type.clone(),
140
- temperature,
141
- max_new_tokens,
142
- top_p,
143
- &dataset_db,
144
- &genedb,
145
- &ai_json,
146
- testing,
147
- )
148
- .await;
149
- let mut llm_json_value: super::super::SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
150
- let mut expected_json_value: super::super::SummaryType = serde_json::from_str(&ques_ans.answer).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
151
- assert_eq!(
152
- llm_json_value.sort_summarytype_struct(),
153
- expected_json_value.sort_summarytype_struct()
154
- );
155
- } else {
156
- panic!("The user input is empty");
137
+ if chart.r#type == "Summary" {
138
+ for ques_ans in chart.TestData {
139
+ let user_input = ques_ans.question;
140
+ if user_input.len() > 0 {
141
+ let llm_output = run_pipeline(
142
+ &user_input,
143
+ comp_model.clone(),
144
+ embedding_model.clone(),
145
+ llm_backend_type.clone(),
146
+ temperature,
147
+ max_new_tokens,
148
+ top_p,
149
+ &dataset_db,
150
+ &genedb,
151
+ &ai_json,
152
+ &airoute,
153
+ testing,
154
+ )
155
+ .await;
156
+ //println!("user_input:{}", user_input);
157
+ //println!("llm_answer:{:?}", llm_output);
158
+ //println!("expected answer:{:?}", &ques_ans.answer);
159
+ let llm_json_value: SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
160
+ //println!(
161
+ // "llm_answer:{:?}",
162
+ // llm_json_value.clone().sort_summarytype_struct()
163
+ //);
164
+ //println!(
165
+ // "expected answer:{:?}",
166
+ // &expected_json_value.clone().sort_summarytype_struct()
167
+ //);
168
+ match ques_ans.answer {
169
+ AnswerFormat::summary_type(sum) => {
170
+ //println!("expected answer:{:?}", &sum);
171
+ assert_eq!(
172
+ llm_json_value.sort_summarytype_struct(),
173
+ sum.sort_summarytype_struct()
174
+ );
175
+ }
176
+ AnswerFormat::DE_type(_) => {
177
+ panic!("DE type not valid for summary")
178
+ }
157
179
  }
180
+ } else {
181
+ panic!("The user input is empty");
158
182
  }
159
183
  }
160
- super::super::Charts::DE(_testdata) => {} // To do
161
184
  }
162
185
  }
163
186
  }