LLM-model-cards / database.py
Blair Yang
New interface
4fe4082
raw
history blame
1.98 kB
from dataclasses import dataclass
from typing import Any
import psycopg2 as pg
@dataclass
class Entry:
model: str
card_file: str
topic: str
question: str
guess: bool
ground_truth: bool
reasoning: str
confidence: int
class PostgreSQL:
hostname: str
dbname: str
user: str
password: str
port: int
# pg.connection
_connection: Any
def __init__(self, hostname: str = "129.153.49.94", port: int = 5432,
dbname: str = "llm-eval", user: str = 'dev', password: str = "llmeval"):
self.hostname = hostname
self.port = port
self.dbname = dbname
self.user = user
self.password = password
self._connection = self._connect()
def insert(self, entry: Entry):
conn = self.get_connection()
with conn.cursor() as cur:
cur.execute(
"""
INSERT INTO responses (model, card_file, topic, question, guess, ground_truth, reasoning, confidence)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
""",
(entry.model, entry.card_file, entry.topic, entry.question, entry.guess, entry.ground_truth,
entry.reasoning, entry.confidence)
)
conn.commit()
def get_connection(self) -> Any:
"""Get the connection to the PostgreSQL database server"""
if self._connection.closed:
self._connection = self._connect()
return self._connection
def _connect(self) -> Any:
"""Connect to the PostgreSQL database server"""
conn = pg.connect(host=self.hostname, port=self.port, dbname=self.dbname, user=self.user,
password=self.password)
with conn.cursor() as cur:
cur.execute("""SET search_path TO %s;""", ("website",))
return conn
def close(self):
# self._connection: pg.connection
self._connection.close()