File size: 1,474 Bytes
02c783d |
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 |
from typing import List
from tenacity import retry, stop_after_attempt, wait_random_exponential
import requests
from models.Base import BaseModel
class GeminiModel(BaseModel):
def __init__(self,
model_id="gemini-2.5-pro-preview-05-06",
api_key=None):
assert api_key is not None, "no api key is provided."
self.model_id = model_id
self.SERVER = "https://llm-api.amd.com/vertex/gemini"
self.HEADERS = {"Ocp-Apim-Subscription-Key": api_key}
@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5))
def generate(self,
messages: List,
temperature=1.0,
presence_penalty=0,
frequency_penalty=0,
max_tokens=30000) -> str:
body = {
"messages": messages,
"max_tokens": max_tokens,
"temperature":temperature,
"top_P": 0.95,
"presence_Penalty": presence_penalty,
"frequency_Penalty": frequency_penalty,
}
response_gemine = requests.post(url=f"{self.SERVER}/{self.model_id}/chat",
json=body,
headers=self.HEADERS)
assert response_gemine.status_code == 200
code_chat_completion_result = response_gemine.json()
return code_chat_completion_result['candidates'][0]['content']['parts'][0]['text'] |