Create db.py

db.py

@@ -0,0 +1,89 @@
+import gradio as gr
+import torch
+import time
+import spaces
+import sqlite3
+import os
+import json
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# --- Database Setup ---
+# Use the persistent storage location in Hugging Face Spaces
+DB_PATH = os.path.join("/data", "model_tests.db")
+
+def init_db():
+    """Initialize the SQLite database with required tables if they don't exist"""
+    with sqlite3.connect(DB_PATH) as conn:
+        cursor = conn.cursor()
+        
+        # Create table for model inputs and outputs
+        cursor.execute('''
+        CREATE TABLE IF NOT EXISTS model_tests (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            analysis_mode TEXT NOT NULL,
+            system_prompt TEXT NOT NULL,
+            input_content TEXT,
+            model_response TEXT NOT NULL,
+            generation_time REAL NOT NULL,
+            tokens_generated INTEGER NOT NULL,
+            temperature REAL NOT NULL,
+            top_p REAL NOT NULL,
+            max_length INTEGER NOT NULL,
+            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
+        )
+        ''')
+        
+        conn.commit()
+    print("Test database initialized successfully")
+
+# Initialize the database when the app starts
+init_db()
+
+# Function to save test results
+def save_test_result(analysis_mode, system_prompt, input_content, model_response, 
+                    generation_time, tokens_generated, temperature, top_p, max_length):
+    with sqlite3.connect(DB_PATH) as conn:
+        cursor = conn.cursor()
+        
+        cursor.execute('''
+        INSERT INTO model_tests 
+        (analysis_mode, system_prompt, input_content, model_response, 
+        generation_time, tokens_generated, temperature, top_p, max_length)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+        ''', (analysis_mode, system_prompt, input_content, model_response, 
+             generation_time, tokens_generated, temperature, top_p, max_length))
+        
+        test_id = cursor.lastrowid
+        conn.commit()
+    
+    return test_id
+
+# Function to retrieve test history
+def get_test_history(limit=10):
+    with sqlite3.connect(DB_PATH) as conn:
+        cursor = conn.cursor()
+        cursor.execute('''
+        SELECT id, analysis_mode, timestamp, generation_time, tokens_generated 
+        FROM model_tests 
+        ORDER BY id DESC LIMIT ?
+        ''', (limit,))
+        history = cursor.fetchall()
+    
+    return history
+
+# Function to get test details by ID
+def get_test_details(test_id):
+    with sqlite3.connect(DB_PATH) as conn:
+        cursor = conn.cursor()
+        cursor.execute('''
+        SELECT * FROM model_tests WHERE id = ?
+        ''', (test_id,))
+        test = cursor.fetchone()
+        
+        if test:
+            # Convert to dictionary
+            columns = [col[0] for col in cursor.description]
+            test_dict = {columns[i]: test[i] for i in range(len(columns))}
+            return test_dict
+        else:
+            return None