Spaces:
Sleeping
Sleeping
File size: 2,567 Bytes
926675f |
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 |
from nodes.LLMNode import *
import time
from utils.util import *
class IO:
def __init__(self, fewshot="\n", model_name="text-davinci-003"):
self.fewshot = fewshot
self.model_name = model_name
self.llm = LLMNode("CoT", model_name, input_type=str, output_type=str)
self.context_prompt = "Answer following questions. Respond directly with no extra words.\n"
self.token_unit_price = get_token_unit_price(model_name)
def run(self, input):
result = {}
st = time.time()
prompt = self.context_prompt + self.fewshot + input + '\n'
response = self.llm.run(prompt, log=True)
result["wall_time"] = time.time() - st
result["input"] = response["input"]
result["output"] = response["output"]
result["prompt_tokens"] = response["prompt_tokens"]
result["completion_tokens"] = response["completion_tokens"]
result["total_tokens"] = response["prompt_tokens"] + response["completion_tokens"]
result["token_cost"] = result["total_tokens"] * self.token_unit_price
result["tool_cost"] = 0
result["total_cost"] = result["token_cost"] + result["tool_cost"]
result["steps"] = 1
return result
class CoT:
def __init__(self, fewshot="\n", model_name="text-davinci-003"):
self.fewshot = fewshot
self.model_name = model_name
self.llm = LLMNode("CoT", model_name, input_type=str, output_type=str)
self.context_prompt = "Answer following questions. Let's think step by step. Give your reasoning process, and then answer the " \
"question in a new line directly with no extra words.\n"
self.token_unit_price = get_token_unit_price(model_name)
def run(self, input):
result = {}
st = time.time()
prompt = self.context_prompt + self.fewshot + input + '\n'
response = self.llm.run(prompt, log=True)
result["wall_time"] = time.time() - st
result["input"] = response["input"]
result["output"] = response["output"]
result["prompt_tokens"] = response["prompt_tokens"]
result["completion_tokens"] = response["completion_tokens"]
result["total_tokens"] = response["prompt_tokens"] + response["completion_tokens"]
result["token_cost"] = result["total_tokens"] * self.token_unit_price
result["tool_cost"] = 0
result["total_cost"] = result["token_cost"] + result["tool_cost"]
result["steps"] = response["output"].count("Step")
return result
|