Spaces:
Runtime error
Runtime error
File size: 4,152 Bytes
36131f5 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | import json
import datetime
import numpy as np
from sqlalchemy import (create_engine, Column, Integer, String, DateTime, Text, Float)
from sqlalchemy.orm import declarative_base, sessionmaker
from sqlalchemy.exc import NoResultFound
import os
Base = declarative_base()
DB_FILE = os.path.join(os.path.dirname(__file__), '..', 'artwork_bandit.db')
ENGINE = create_engine(f'sqlite:///{DB_FILE}', connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(bind=ENGINE)
class Impression(Base):
__tablename__ = 'impressions'
id = Column(Integer, primary_key=True, index=True, autoincrement=True)
user_id = Column(String, index=True)
content_id = Column(String, index=True)
artwork_id = Column(String, index=True)
timestamp = Column(DateTime, default=datetime.datetime.utcnow)
reward = Column(Float, nullable=True)
class BanditState(Base):
__tablename__ = 'bandit_state'
id = Column(Integer, primary_key=True)
arm_id = Column(String, unique=True, index=True)
A_matrix = Column(Text)
b_vector = Column(Text)
updated_at = Column(DateTime, default=datetime.datetime.utcnow)
class FeatureCache(Base):
__tablename__ = 'feature_cache'
id = Column(Integer, primary_key=True)
entity_type = Column(String, index=True)
entity_id = Column(String, index=True)
embedding = Column(Text)
created_at = Column(DateTime, default=datetime.datetime.utcnow)
def init_db():
Base.metadata.create_all(bind=ENGINE)
def log_impression(user_id: str, content_id: str, artwork_id: str):
db = SessionLocal()
imp = Impression(user_id=user_id, content_id=content_id, artwork_id=artwork_id)
db.add(imp)
db.commit()
db.refresh(imp)
db.close()
return imp.id
def log_reward(impression_id: int, reward: float):
db = SessionLocal()
imp = db.query(Impression).filter(Impression.id == impression_id).one()
imp.reward = float(reward)
db.commit()
db.close()
def save_bandit_arm(arm_id: str, A: np.ndarray, b: np.ndarray):
db = SessionLocal()
textA = json.dumps(A.tolist())
textb = json.dumps(b.tolist())
try:
state = db.query(BanditState).filter(BanditState.arm_id == arm_id).one()
state.A_matrix = textA
state.b_vector = textb
state.updated_at = datetime.datetime.utcnow()
except NoResultFound:
state = BanditState(arm_id=arm_id, A_matrix=textA, b_vector=textb)
db.add(state)
db.commit()
db.close()
def load_bandit_arm(arm_id: str):
db = SessionLocal()
try:
state = db.query(BanditState).filter(BanditState.arm_id == arm_id).one()
A = np.array(json.loads(state.A_matrix))
b = np.array(json.loads(state.b_vector))
db.close()
return {'arm_id': arm_id, 'A': A, 'b': b}
except NoResultFound:
db.close()
return None
def save_embedding(entity_type: str, entity_id: str, embedding: np.ndarray):
db = SessionLocal()
text = json.dumps(embedding.tolist())
try:
existing = db.query(FeatureCache).filter(FeatureCache.entity_type==entity_type, FeatureCache.entity_id==entity_id).one()
existing.embedding = text
except NoResultFound:
entry = FeatureCache(entity_type=entity_type, entity_id=entity_id, embedding=text)
db.add(entry)
db.commit()
db.close()
def load_embedding(entity_type: str, entity_id: str):
db = SessionLocal()
try:
entry = db.query(FeatureCache).filter(FeatureCache.entity_type==entity_type, FeatureCache.entity_id==entity_id).one()
emb = np.array(json.loads(entry.embedding))
db.close()
return emb
except NoResultFound:
db.close()
return None
def get_impression(impression_id: int):
db = SessionLocal()
try:
imp = db.query(Impression).filter(Impression.id==impression_id).one()
db.expunge(imp)
db.close()
return imp
except NoResultFound:
db.close()
return None
def impressions_with_rewards():
db = SessionLocal()
rows = db.query(Impression).filter(Impression.reward != None).all()
db.close()
return rows
|