Spaces:
Running
Running
File size: 1,978 Bytes
4fe4082 |
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 56 57 58 59 60 61 62 63 64 65 66 67 |
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()
|