forked from A-I-nstein/YouTube-TLDR
-
Notifications
You must be signed in to change notification settings - Fork 0
/
palm_api.py
36 lines (29 loc) · 1012 Bytes
/
palm_api.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
import os
import google.generativeai as palm
from langchain.llms.base import LLM
from typing import Any, List, Mapping, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun
palm.configure(api_key=os.getenv('PALM_API_KEY'))
class PALM(LLM):
@property
def _llm_type(self) -> str:
return "PALM API"
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
if stop is not None:
raise ValueError("stop kwargs are not permitted.")
models = [m for m in palm.list_models() if 'generateText' in m.supported_generation_methods]
model = models[0].name
completion = palm.generate_text(
model=model,
prompt=prompt,
temperature=0,
# The maximum length of the response
max_output_tokens=800,
)
return completion.result