File size: 1,877 Bytes
2b659a0 f404d37 2b659a0 f404d37 2b659a0 |
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 |
"""Load models to use them as a narrator and a common-sense oracle in the PAYADOR pipeline."""
import google.generativeai as genai
import requests
import os
class GeminiModel():
def __init__ (self, api_key_file:str, model_name:str = "gemini-pro") -> None:
""""Initialize the Gemini model using an API key."""
self.safety_settings = [
{
"category": "HARM_CATEGORY_DANGEROUS",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_HARASSMENT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_NONE",
},
]
genai.configure(api_key=os.getenv(api_key_file))
self.model = genai.GenerativeModel(model_name)
def prompt_model(self,prompt: str) -> str:
"""Prompt the Gemini model."""
return self.model.generate_content(prompt, safety_settings=self.safety_settings).text
def prompt_HF_API (prompt: str, model: str = "microsoft/Phi-3-mini-4k-instruct", api_key_file: str = "HF_API_key"):
API_URL = f"https://api-inference.huggingface.co/models/{model}"
headers = {"Authorization": f"Bearer {get_api_key(api_key_file)}"}
payload = {"inputs": prompt}
output = requests.post(API_URL, headers=headers, json=payload).json()
return output[0]["generated_text"]
def get_api_key(path: str) -> str:
"""Load an API key from path."""
key = ""
with open(path) as f:
key = f.readline()
return key |