File size: 4,034 Bytes
681ffab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e88e1f
 
681ffab
 
 
 
 
 
 
 
 
 
 
 
0e88e1f
 
681ffab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e88e1f
681ffab
 
 
 
 
 
 
 
 
 
6dce795
 
 
 
 
 
 
 
 
 
 
681ffab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5d73a34
681ffab
 
 
 
 
 
 
 
 
0e957b2
681ffab
 
 
 
 
 
0e957b2
681ffab
 
 
 
 
 
 
0e957b2
681ffab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e957b2
 
5d73a34
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
131
132
133
134
135
136
137
138
139
140
141
142
143
# 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
        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(
                    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
        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()

    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}")

    model = openai.ChatCompletion.create(
        model=model_id,
        messages=[],
        temperature=0.0,
        max_tokens=0,
        n=1,
        stop=None,
        log_level="info",
    )

    return model


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("\n")

    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


# 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)