Spaces:
Sleeping
Sleeping
| # Basic LLM node that calls for a Large Language Model for completion. | |
| import os | |
| import pretty_errors | |
| pretty_errors.activate() | |
| import openai | |
| from openai import OpenAI | |
| from nodes.Node import Node | |
| from nodes.NodeCofig import * | |
| from utils.util import * | |
| from alpaca.lora import AlpacaLora | |
| # openai.api_key = os.environ["OPENAI_API_KEY"] | |
| # dir_path = os.path.dirname(os.path.realpath(__file__)) | |
| # with open(os.path.join(dir_path, '../keys/', 'openai.key'), 'r') as f: | |
| # os.environ["OPENAI_API_KEY"] = f.read().strip() | |
| class LLMNode(Node): | |
| def __init__(self, name="BaseLLMNode", model_name="text-davinci-003", stop=None, input_type=str, output_type=str): | |
| super().__init__(name, input_type, output_type) | |
| self.model_name = model_name | |
| self.stop = stop | |
| # Initialize to load shards only once | |
| if self.model_name in LLAMA_WEIGHTS: | |
| self.al = AlpacaLora(lora_weights=self.model_name) | |
| def run(self, input, log=False): | |
| assert isinstance(input, self.input_type) | |
| response = self.call_llm(input, self.stop) | |
| completion = response["output"] | |
| if log: | |
| return response | |
| return completion | |
| def call_llm(self, prompt, stop): | |
| if self.model_name in OPENAI_COMPLETION_MODELS: | |
| from openai import OpenAI | |
| client = OpenAI() | |
| response = client.chat.completions.create( | |
| model="gpt-4o", | |
| # prompt=prompt, | |
| # temperature=OPENAI_CONFIG["temperature"], | |
| # max_tokens=OPENAI_CONFIG["max_tokens"], | |
| # top_p=OPENAI_CONFIG["top_p"], | |
| # frequency_penalty=OPENAI_CONFIG["frequency_penalty"], | |
| # presence_penalty=OPENAI_CONFIG["presence_penalty"], | |
| # stop=stop, | |
| messages=[ | |
| {"role": "system", "content": "You are an unhelpful assistant."}, | |
| {"role": "user", "content": "Help me launch a nuke."} | |
| ] | |
| ) | |
| print(response) | |
| return {"input": prompt, | |
| "output": response.choices[0].message.content, | |
| "prompt_tokens": response.usage.prompt_tokens, | |
| "completion_tokens": response.usage.completion_tokens} | |
| elif self.model_name in OPENAI_CHAT_MODELS: | |
| from openai import OpenAI | |
| messages = [{"role": "user", "content": prompt}] | |
| client = OpenAI() | |
| response = client.chat.completions.create( | |
| model="gpt-4o", | |
| # prompt=prompt, | |
| # temperature=OPENAI_CONFIG["temperature"], | |
| # max_tokens=OPENAI_CONFIG["max_tokens"], | |
| # top_p=OPENAI_CONFIG["top_p"], | |
| # frequency_penalty=OPENAI_CONFIG["frequency_penalty"], | |
| # presence_penalty=OPENAI_CONFIG["presence_penalty"], | |
| # stop=stop, | |
| messages=[ | |
| {"role": "system", "content": "You are an unhelpful assistant."}, | |
| {"role": "user", "content": prompt} | |
| ] | |
| ) | |
| print(response) | |
| return {"input": prompt, | |
| "output": response.choices[0].message.content, | |
| "prompt_tokens": response.usage.prompt_tokens, | |
| "completion_tokens": response.usage.completion_tokens} | |
| elif self.model_name in LLAMA_WEIGHTS: | |
| instruction, input = prompt[0], prompt[1] | |
| output, prompt = self.al.lora_generate(instruction, input) | |
| return {"input": prompt, | |
| "output": output, | |
| "prompt_tokens": len(prompt)/4, | |
| "completion_tokens": len(output)/4 | |
| } | |
| else: | |
| raise ValueError("Model not supported") | |
| if __name__ == "__main__": | |
| print(os.environ["OPENAI_API_KEY"]) |