File size: 3,682 Bytes
681ffab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e957b2
681ffab
 
 
 
 
 
 
 
 
0e957b2
681ffab
 
 
 
 
 
0e957b2
681ffab
 
 
 
 
 
 
0e957b2
681ffab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e957b2
 
 
 
 
 
 
 
 
 
 
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
# custom text generation llm classes

import warnings
import logging
import os

import openai

# supress warnings
warnings.filterwarnings("ignore")


class llm_boiler:
    def __init__(self, model_id, openai_key):
        self.model_id = model_id
        self.openai_key = openai_key
        for f_idx, run_function in enumerate(MODEL_FUNCTIONS):
            if run_function.__name__.lower() in self.model_id:
                print(
                    f"Load function recognized for {self.model_id}: {LOAD_MODEL_FUNCTIONS[f_idx].__name__}"
                )
                self.load_fn = LOAD_MODEL_FUNCTIONS[f_idx]
        for run_function in MODEL_FUNCTIONS:
            if run_function.__name__.lower() in self.model_id:
                print(
                    f"Run function recognized for {self.model_id}: {run_function.__name__.lower()}"
                )
                self.run_fn = run_function
        self.model = self.load_fn(self.model_id, self.openai_key)
        self.name = self.run_fn.__name__.lower()

    def run(
        self,
        prompt,
        temperature,
    ):
        return self.run_fn(
            model=self.model,
            prompt=prompt,
            temperature=temperature,
        )


LOAD_MODEL_FUNCTIONS = []
MODEL_FUNCTIONS = []


# gpt models
def gpt_loader(model_id: str, openai_key: str):
    # Load your API key from an environment variable or secret management service
    openai.api_key = openai_key  # os.getenv("OPENAI_API_KEY")
    logging.warning(f"model id: {model_id}")

    return model_id


LOAD_MODEL_FUNCTIONS.append(gpt_loader)


def gpt(
    model: str,
    prompt: str,
    temperature: int,
) -> str:
    """
    Initialize the pipeline
    Uses Hugging Face GenerationConfig defaults
        https://huggingface.co./docs/transformers/v4.29.1/en/main_classes/text_generation#transformers.GenerationConfig
    Args:
        model (str): openai model key
        tokenizer (str): openai model key
        prompt (str): Prompt for text generation
        max_new_tokens (int, optional): Max new tokens after the prompt to generate. Defaults to 128.
        temperature (float, optional): The value used to modulate the next token probabilities.
            Defaults to 1.0
    """
    conversation = prompt.split("")

    messages = []
    for turn in conversation:
        first_word = turn.split("\n")[0]

        if first_word == "system":
            messages.append(
                {
                    "role": "system",
                    "content": turn.replace("system\n", "").replace("\n", ""),
                }
            )
        elif first_word == "user":
            messages.append(
                {
                    "role": "user",
                    "content": turn.replace("user\n", "").replace("\n", ""),
                }
            )
        elif first_word == "assistant":
            messages.append(
                {
                    "role": "assistant",
                    "content": turn.replace("assistant\n", "").replace(
                        "\n", ""
                    ),
                }
            )

    logging.warning(f"Input to openai api call: {messages}")

    chat_completion = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=temperature,
        stream=True,
    )
    return chat_completion


MODEL_FUNCTIONS.append(gpt)

# Define the model and its parameters
model_id = "dfurman/chat-gpt-3.5-turbo"
openai_key = "<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>"

model = llm_boiler(model_id, openai_key)

prompt = "Hello, how are you?"
temperature = 0.8

response = model.run(prompt, temperature)