update test file

llamatest.py

@@ -1,15 +1,22 @@
-import transformers
+from transformers import pipeline
 import torch
 import os
 from tot.models import gpt
 
 os.environ["HF_TOKEN"] = os.getenv("HUGGINGTOKEN")
 
-# model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+
 
 def testThoughtGeneration():
+    model_id = "meta-llama/Meta-Llama-3.1-70B-Instruct"
+    model_pipeline = pipeline(
+        "text-generation",
+        model=model_id,
+        model_kwargs={"torch_dtype": torch.bfloat16},
+        device_map="auto",
+    )
     my_propose_prompt20 = 'Input: 2 8 8 14\nPossible next steps:\n2 + 8 = 10 (left: 8 10 14)\n8 / 2 = 4 (left: 4 8 14)\n14 + 2 = 16 (left: 8 8 16)\n2 * 8 = 16 (left: 8 14 16)\n8 - 2 = 6 (left: 6 8 14)\n14 - 8 = 6 (left: 2 6 8)\n14 / 2 = 7 (left: 7 8 8)\n14 - 2 = 12 (left: 8 8 12)\nInput: 4 5 6 10\nPossible next steps:\n'
-    result = gpt(my_propose_prompt20, max_tokens=1000, n=1, stop=["Input"])
+    result = gpt(my_propose_prompt20, model_pipeline, max_tokens=1000, n=1, stop=["Input"])
     print(result)
 
 testThoughtGeneration()