lstetson commited on
Commit
d80714e
1 Parent(s): 61e6bcc

Add fullset to main test

Browse files
Files changed (1) hide show
  1. tests/test_main.py +21 -2
tests/test_main.py CHANGED
@@ -21,6 +21,23 @@ def test_with_subset():
21
  question = "How should I train for anerobic capacity?"
22
  print("Question: ", question)
23
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
  answer = run_query(question,
25
  db_path=db_path,
26
  num_rel_segments=10,
@@ -32,11 +49,13 @@ def test_with_subset():
32
 
33
 
34
  if __name__ == '__main__':
35
- choice = input("Enter 1 for test_with_subset, 2 for test_with_llm_video: ")
36
 
37
  if choice == "1":
38
- test_with_subset()
39
  elif choice == "2":
 
 
40
  test_with_llm_video()
41
  else:
42
  print("Invalid choice")
 
21
  question = "How should I train for anerobic capacity?"
22
  print("Question: ", question)
23
 
24
+ answer = run_query(question,
25
+ db_path=db_path,
26
+ num_rel_segments=10,
27
+ llm_model="gpt-3.5-turbo-0125",
28
+ llm_temp=0.1
29
+ )
30
+
31
+ print("Answer: ", answer)
32
+
33
+
34
+ @timeit
35
+ def test_with_fullset():
36
+ db_path = "data/videos.db"
37
+
38
+ question = "How should I train for anerobic capacity?"
39
+ print("Question: ", question)
40
+
41
  answer = run_query(question,
42
  db_path=db_path,
43
  num_rel_segments=10,
 
49
 
50
 
51
  if __name__ == '__main__':
52
+ choice = input("Test with:\n1. fullset\n2. subset\n3. LLM video\nEnter option number: ")
53
 
54
  if choice == "1":
55
+ test_with_fullset()
56
  elif choice == "2":
57
+ test_with_subset()
58
+ elif choice == "3":
59
  test_with_llm_video()
60
  else:
61
  print("Invalid choice")