-
Notifications
You must be signed in to change notification settings - Fork 0
/
genai.py
38 lines (31 loc) · 969 Bytes
/
genai.py
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
import os
from langchain_google_vertexai import ChatVertexAI
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
llm :dict[str, ChatVertexAI|None] = {
"gemini-pro": None,
"gemini-pro-vision": None,
}
memory = ConversationBufferMemory()
_debug: bool = 'DEBUG' in os.environ
def _init(params: dict = {}):
default_params: dict = {
"model_name": "",
"temperature": 0.6,
"max_output_tokens": 1024,
"top_p": 0.8,
"top_k": 40,
}
params.update(default_params)
return params
def get_llm(kind: str = 'gemini-pro') -> ChatVertexAI:
global llm, memory
if kind in llm and llm[kind] is None:
p = _init()
p["model_name"] = kind
llm[kind] = ChatVertexAI(**p)
print(f"Generated LLM:{kind}")
return llm[kind]
def run(text_string: str) -> str:
response = get_llm().predict(input=text_string)
return response