| 114 | traj = [] |
| 115 | |
| 116 | class LLMNode(Node): |
| 117 | def __init__(self, name="BaseLLMNode", model_name="text-davinci-003", stop=None, input_type=str, output_type=str): |
| 118 | super().__init__(name, input_type, output_type) |
| 119 | self.model_name = model_name |
| 120 | self.stop = stop |
| 121 | |
| 122 | # Initialize to load shards only once |
| 123 | if self.model_name in LLAMA_WEIGHTS: |
| 124 | self.al = AlpacaLora(lora_weights=self.model_name) |
| 125 | |
| 126 | def run(self, input, log=False): |
| 127 | assert isinstance(input, self.input_type) |
| 128 | response = self.call_llm(input, self.stop) |
| 129 | completion = response["output"] |
| 130 | if log: |
| 131 | return response |
| 132 | return completion |
| 133 | |
| 134 | def call_llm(self, prompt, stop): |
| 135 | if self.model_name in OPENAI_COMPLETION_MODELS: |
| 136 | response = openai.Completion.create( |
| 137 | model=self.model_name, |
| 138 | prompt=prompt, |
| 139 | temperature=OPENAI_CONFIG["temperature"], |
| 140 | max_tokens=OPENAI_CONFIG["max_tokens"], |
| 141 | top_p=OPENAI_CONFIG["top_p"], |
| 142 | frequency_penalty=OPENAI_CONFIG["frequency_penalty"], |
| 143 | presence_penalty=OPENAI_CONFIG["presence_penalty"], |
| 144 | stop=stop |
| 145 | ) |
| 146 | return {"input": prompt, |
| 147 | "output": response["choices"][0]["text"], |
| 148 | "prompt_tokens": response["usage"]["prompt_tokens"], |
| 149 | "completion_tokens": response["usage"]["completion_tokens"]} |
| 150 | elif self.model_name in OPENAI_CHAT_MODELS: |
| 151 | print('*****GPT-4*****') |
| 152 | messages = [{"role": "user", "content": prompt}] |
| 153 | while True: |
| 154 | try: |
| 155 | response = openai.ChatCompletion.create( |
| 156 | model=self.model_name, |
| 157 | messages=messages, |
| 158 | # temperature=OPENAI_CONFIG["temperature"], |
| 159 | temperature=0.0, |
| 160 | max_tokens=OPENAI_CONFIG["max_tokens"], |
| 161 | # top_p=OPENAI_CONFIG["top_p"], |
| 162 | # frequency_penalty=OPENAI_CONFIG["frequency_penalty"], |
| 163 | # presence_penalty=OPENAI_CONFIG["presence_penalty"], |
| 164 | # stop=stop |
| 165 | ) |
| 166 | break |
| 167 | except: |
| 168 | continue |
| 169 | |
| 170 | traj.append({ |
| 171 | 'id': f'{ID}_{len(traj)}', |
| 172 | 'conversations': [ |
| 173 | { |