Original implementation

example.py

@@ -28,14 +28,15 @@
 T = 3
 b = 5
 vth = 0.5
+pruning_threshold = 0.5
 timeout = 10
 confidence = 0.8 #cmodel is confident on performance
 max_iterations = 40 #tree branh nodes 
 convergence_threshold = 0.01
 convergence_count = 5
 
 #call the solve emthod with the input problem and other params
-solution = tree_of_thoughts.solve(input_problem, k, T, b, vth, timeout, confidence, max_iterations, convergence_threshold, convergence_count)
+solution = tree_of_thoughts.solve(input_problem, k, T, b, vth, timeout, pruning_threshold, confidence, max_iterations, convergence_threshold, convergence_count)
 
 
 

huggingfaceExample.py

@@ -1,4 +1,5 @@
-from tree_of_thoughts.treeofthoughts import OpenAILanguageModel, GuidanceOpenAILanguageModel, TreeofThoughts, OptimizedOpenAILanguageModel, OptimizedTreeofThoughts, HuggingLanguageModel
+from tree_of_thoughts.treeofthoughts import TreeofThoughts, OptimizedTreeofThoughts
+from tree_of_thoughts.huggingModels import HuggingLanguageModel
 # from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # model_name="gpt2"
@@ -18,9 +19,10 @@
 #         return generated_text
 
 # Initialize the HuggingLanguageModel with the GPT-2 model
-model_name = "gpt2"
+model_name = r"C:\Users\David\.cache\huggingface\modules\transformers_modules\chatglm6b"
+tokenizer_name = r"C:\Users\David\.cache\huggingface\modules\transformers_modules\chatglm6b"
 model = HuggingLanguageModel(model_name, 
-                             model_Tokenizer="gpt2", 
+                             model_tokenizer=tokenizer_name,
                              verbose=True)
 
 
@@ -35,7 +37,6 @@
 # value or vote
 evaluation_strategy = "value"
 
-
 gpt2_model = HuggingLanguageModel(model_name)
 
 tree_of_thoughts= OptimizedTreeofThoughts(model, search_algorithm)
@@ -50,9 +51,12 @@
 max_iterations = 40 #tree branh nodes 
 convergence_threshold = 0.01
 convergence_count = 5
+pruning_threshold = 0.3
 
 #call the solve emthod with the input problem and other params
-solution = tree_of_thoughts.solve(input_problem, k, T, b, vth, timeout, confidence, max_iterations, convergence_threshold, convergence_count)
+solution = tree_of_thoughts.solve(input_problem, k, T, b, vth, timeout,
+                                  confidence_threshold=confidence, max_iterations=max_iterations,
+                                  convergence_threshold=convergence_threshold, convergence_count=convergence_count)
 
 #use the solution in your production environment
 print(f"solution: {solution}")

tree_of_thoughts/huggingModels.py

@@ -1,12 +1,12 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
 from transformers import pipeline
 from .abstractLanguageModel import AbstractLanguageModel
 
 
 class HuggingLanguageModel(AbstractLanguageModel):
     def __init__(self, model_name, model_tokenizer=None, verbose=False):
-        self.model = AutoModelForCausalLM.from_pretrained(model_name)
-        self.tokenizer = AutoTokenizer.from_pretrained(model_tokenizer or model_name)
+        self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to('cpu', dtype=float)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_tokenizer or model_name, trust_remote_code=True)
         self.verbose = verbose
 
     def generate_thoughts(self, state, k, max_length=100):
@@ -28,7 +28,7 @@ def generate_thoughts(self, state, k, max_length=100):
 
         return thoughts
 
-    def evaluate_states(self, states, inital_prompt, max_length=10):
+    def evaluate_states(self, states, initial_prompt, max_length=10):
         state_values = {}
         for state in states:
             state_text = ' '.join(state)