Update src/llm_boilers.py

src/llm_boilers.py

@@ -14,8 +14,6 @@ class llm_boiler:
     def __init__(self, model_id, openai_key):
         self.model_id = model_id
         self.openai_key = openai_key
-        self.load_fn = None  # Add the load_fn attribute
-        self.run_fn = None  # Add the run_fn attribute
         for f_idx, run_function in enumerate(MODEL_FUNCTIONS):
             if run_function.__name__.lower() in self.model_id:
                 print(
@@ -28,8 +26,6 @@ class llm_boiler:
                     f"Run function recognized for {self.model_id}: {run_function.__name__.lower()}"
                 )
                 self.run_fn = run_function
-        if self.load_fn is None or self.run_fn is None:
-            raise ValueError("Invalid model_id")
         self.model = self.load_fn(self.model_id, self.openai_key)
         self.name = self.run_fn.__name__.lower()
 
@@ -45,7 +41,6 @@ class llm_boiler:
         )
 
 
-
 LOAD_MODEL_FUNCTIONS = []
 MODEL_FUNCTIONS = []
 
@@ -56,17 +51,7 @@ def gpt_loader(model_id: str, openai_key: str):
     openai.api_key = openai_key  # os.getenv("OPENAI_API_KEY")
     logging.warning(f"model id: {model_id}")
 
-    model = openai.ChatCompletion.create(
-        model=model_id,
-        messages=[],
-        temperature=0.0,
-        max_tokens=0,
-        n=1,
-        stop=None,
-        log_level="info",
-    )
-
-    return model
+    return model_id
 
 
 LOAD_MODEL_FUNCTIONS.append(gpt_loader)
@@ -89,7 +74,7 @@ def gpt(
         temperature (float, optional): The value used to modulate the next token probabilities.
             Defaults to 1.0
     """
-    conversation = prompt.split("\n")
+    conversation = prompt.split("<|im_start|>")
 
     messages = []
     for turn in conversation:
@@ -99,14 +84,14 @@ def gpt(
             messages.append(
                 {
                     "role": "system",
-                    "content": turn.replace("system\n", "").replace("\n", ""),
+                    "content": turn.replace("system\n", "").replace("<|im_end|>\n", ""),
                 }
             )
         elif first_word == "user":
             messages.append(
                 {
                     "role": "user",
-                    "content": turn.replace("user\n", "").replace("\n", ""),
+                    "content": turn.replace("user\n", "").replace("<|im_end|>\n", ""),
                 }
             )
         elif first_word == "assistant":
@@ -114,7 +99,7 @@ def gpt(
                 {
                     "role": "assistant",
                     "content": turn.replace("assistant\n", "").replace(
-                        "\n", ""
+                        "<|im_end|>\n", ""
                     ),
                 }
             )
@@ -130,13 +115,4 @@ def gpt(
     return chat_completion
 
 
-# Define the model and its parameters
-model_id = "dfurman/chat-gpt-3.5-turbo"
-openai_key = os.getenv("API_KEY")
-
-model = llm_boiler(model_id, openai_key)
-
-prompt = "Hello, how are you?"
-temperature = 0.8
-
-response = model.run(prompt, temperature)
+MODEL_FUNCTIONS.append(gpt)