Update app.py

app.py

@@ -1,36 +1,29 @@
 import gradio as gr
 import json
-import os
 import hashlib
-import tempfile
+import sqlite3
+import os
 
 class JSONLViewer:
-    def __init__(self, data_file_path):
+    def __init__(self, data_file_path, db_path):
         self.data_file_path = data_file_path
-        self.state_file_path = os.path.join(tempfile.gettempdir(), 'app_state.json')
+        self.db_path = db_path
         self.current_index = 0
         self.data = []
-        self.state = {}
         self.load_data()
-        self.load_state()
+        self.init_db()
 
     def load_data(self):
-        if os.path.exists(self.data_file_path):
-            with open(self.data_file_path, 'r', encoding='utf-8') as file:
-                self.data = [json.loads(line) for line in file]
-        else:
-            self.data = []
-
-    def load_state(self):
-        if os.path.exists(self.state_file_path):
-            with open(self.state_file_path, 'r', encoding='utf-8') as file:
-                self.state = json.load(file)
-        else:
-            self.state = {}
-
-    def save_state(self):
-        with open(self.state_file_path, 'w', encoding='utf-8') as file:
-            json.dump(self.state, file, ensure_ascii=False, indent=2)
+        with open(self.data_file_path, 'r', encoding='utf-8') as file:
+            self.data = [json.loads(line) for line in file]
+
+    def init_db(self):
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        c.execute('''CREATE TABLE IF NOT EXISTS states
+                     (record_id TEXT PRIMARY KEY, status TEXT)''')
+        conn.commit()
+        conn.close()
 
     def get_current_record(self):
         if 0 <= self.current_index < len(self.data):
@@ -41,16 +34,21 @@ class JSONLViewer:
         record_str = json.dumps(record, sort_keys=True)
         return hashlib.md5(record_str.encode()).hexdigest()
 
-    def update_status(self, status):
-        current_record = self.get_current_record()
-        if current_record:
-            record_id = self.get_record_id(current_record)
-            self.state[record_id] = status
-            self.save_state()
-
-    def get_status(self, record):
-        record_id = self.get_record_id(record)
-        return self.state.get(record_id, "")
+    def get_status(self, record_id):
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        c.execute("SELECT status FROM states WHERE record_id = ?", (record_id,))
+        result = c.fetchone()
+        conn.close()
+        return result[0] if result else ""
+
+    def set_status(self, record_id, status):
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        c.execute("INSERT OR REPLACE INTO states (record_id, status) VALUES (?, ?)",
+                  (record_id, status))
+        conn.commit()
+        conn.close()
 
     def move_prev(self):
         if self.current_index > 0:
@@ -62,7 +60,7 @@ class JSONLViewer:
             self.current_index += 1
         return self.get_current_record()
 
-viewer = JSONLViewer('plik.jsonl')
+viewer = JSONLViewer('plik.jsonl', 'states.db')
 
 def update_ui(record):
     if record:
@@ -71,7 +69,8 @@ def update_ui(record):
         rejected = record['rejected']['value'] if 'rejected' in record else ""
         chosen_score = record['chosen_score'] if 'chosen_score' in record else ""
         rejected_score = record['rejected_score'] if 'rejected_score' in record else ""
-        status = viewer.get_status(record)
+        record_id = viewer.get_record_id(record)
+        status = viewer.get_status(record_id)
         return instruction, chosen, rejected, chosen_score, rejected_score, status
     return "", "", "", "", "", ""
 
@@ -84,19 +83,23 @@ def on_next():
     return update_ui(record)
 
 def on_ok():
-    viewer.update_status("ok")
     record = viewer.get_current_record()
+    if record:
+        record_id = viewer.get_record_id(record)
+        viewer.set_status(record_id, "ok")
     return update_ui(record)
 
 def on_rejected():
-    viewer.update_status("rejected")
     record = viewer.get_current_record()
+    if record:
+        record_id = viewer.get_record_id(record)
+        viewer.set_status(record_id, "rejected")
     return update_ui(record)
 
 with gr.Blocks(css="button.ok-button { background-color: #4CAF50 !important; }") as demo:
-    instruction = gr.Textbox(label="INSTRUCTION", lines=2, max_lines=2)
-    chosen = gr.Textbox(label="CHOSEN", lines=12, max_lines=12)
-    rejected = gr.Textbox(label="REJECTED", lines=12, max_lines=12)
+    instruction = gr.Textbox(label="INSTRUCTION", lines=8, max_lines=8)
+    chosen = gr.Textbox(label="CHOSEN", lines=8, max_lines=8)
+    rejected = gr.Textbox(label="REJECTED", lines=8, max_lines=8)
     chosen_score = gr.Number(label="CHOSEN SCORE")
     rejected_score = gr.Number(label="REJECTED SCORE")
     status = gr.Textbox(label="STATUS")