Update app.py

app.py

@@ -6,10 +6,88 @@ import plotly.express as px
 from collections import defaultdict
 from datetime import datetime
 import os
+from datasets import load_dataset
+import sqlite3
 
+def initialize_database():
+    conn = sqlite3.connect('afrimmlu_results.db')
+    cursor = conn.cursor()
+    
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS summary_results (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            language TEXT,
+            subject TEXT,
+            accuracy REAL,
+            timestamp TEXT
+        )
+    ''')
+
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS detailed_results (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            language TEXT,
+            timestamp TEXT,
+            subject TEXT,
+            question TEXT,
+            model_answer TEXT,
+            correct_answer TEXT,
+            is_correct INTEGER,
+            total_tokens INTEGER
+        )
+    ''')
+
+    conn.commit()
+    conn.close()
+
+def save_results_to_database(language, summary_results, detailed_results):
+    conn = sqlite3.connect('afrimmlu_results.db')
+    cursor = conn.cursor()
+    timestamp = datetime.now().isoformat()
+
+    # Save summary results
+    for subject, accuracy in summary_results.items():
+        cursor.execute('''
+            INSERT INTO summary_results (language, subject, accuracy, timestamp)
+            VALUES (?, ?, ?, ?)
+        ''', (language, subject, accuracy, timestamp))
 
+    # Save detailed results
+    for result in detailed_results:
+        cursor.execute('''
+            INSERT INTO detailed_results (
+                language, timestamp, subject, question, model_answer, 
+                correct_answer, is_correct, total_tokens
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+        ''', (
+            language,
+            result['timestamp'],
+            result['subject'],
+            result['question'],
+            result['model_answer'],
+            result['correct_answer'],
+            int(result['is_correct']),
+            result['total_tokens']
+        ))
 
-os.environ['DEEPSEEK_API_KEY']
+    conn.commit()
+    conn.close()
+
+def load_afrimmlu_data(language_code="swa"):
+    """
+    Load AfriMMLU dataset for a specific language.
+    """
+    try:
+        dataset = load_dataset(
+            'masakhane/afrimmlu', 
+            language_code, 
+            token=os.environ['HF_TOKEN'],
+        )
+        test_data = dataset['test'].to_list()
+        return test_data
+    except Exception as e:
+        print(f"Error loading dataset: {str(e)}")
+        return None
 
 def preprocess_dataset(test_data):
     """
@@ -32,7 +110,7 @@ def preprocess_dataset(test_data):
         preprocessed_data.append(example)
     return preprocessed_data
 
-def evaluate_afrimmlu(test_data, model_name="deepseek-chat"):
+def evaluate_afrimmlu(test_data, model_name="deepseek/deepseek-chat", language="swa"):
     """
     Evaluate the model on the AfriMMLU dataset.
     """
@@ -79,7 +157,6 @@ def evaluate_afrimmlu(test_data, model_name="deepseek-chat"):
             total += 1
             subject_results[subject]["total"] += 1
 
-            # Store detailed results
             results.append({
                 'timestamp': datetime.now().isoformat(),
                 'subject': subject,
@@ -94,22 +171,14 @@ def evaluate_afrimmlu(test_data, model_name="deepseek-chat"):
             print(f"Error processing question: {str(e)}")
             continue
 
-    # Calculate accuracies
     accuracy = (correct / total * 100) if total > 0 else 0
     subject_accuracy = {
         subject: (stats["correct"] / stats["total"] * 100) if stats["total"] > 0 else 0
         for subject, stats in subject_results.items()
     }
 
-    # Export results to CSV
-    df = pd.DataFrame(results)
-    df.to_csv('detailed_results.csv', index=False)
-
-    # Export summary to CSV
-    summary_data = [{'subject': subject, 'accuracy': acc} 
-                   for subject, acc in subject_accuracy.items()]
-    summary_data.append({'subject': 'Overall', 'accuracy': accuracy})
-    pd.DataFrame(summary_data).to_csv('summary_results.csv', index=False)
+    # Save results to database
+    save_results_to_database(language, {**subject_accuracy, 'Overall': accuracy}, results)
 
     return {
         "accuracy": accuracy,
@@ -143,41 +212,34 @@ def create_visualization(results_dict):
     
     return summary_df, fig
 
-def evaluate_and_display(test_file, model_name):
-    # Load and preprocess data
-    test_data = pd.read_json(test_file.name)
-    preprocessed_data = preprocess_dataset(test_data.to_dict('records'))
-    
-    # Run evaluation
-    results = evaluate_afrimmlu(preprocessed_data, model_name)
-    
-    # Create visualizations
-    summary_df, plot = create_visualization(results)
-    
-    # Load detailed results with error handling
-    try:
-        detailed_df = pd.read_csv('detailed_results.csv')
-    except (FileNotFoundError, pd.errors.EmptyDataError):
-        detailed_df = pd.DataFrame(results["detailed_results"])
-    
-    return summary_df, plot, detailed_df
-
-
 def create_gradio_interface():
     """
     Create and configure the Gradio interface.
     """
+    language_options = {
+        "swa": "Swahili",
+        "yor": "Yoruba",
+        "wol": "Wolof",
+        "lin": "Lingala",
+        "ewe": "Ewe",
+        "ibo": "Igbo"
+    }
+    
+    # Initialize database
+    initialize_database()
+    
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
         # AfriMMLU Evaluation Dashboard
-        Upload your test data and select a model to evaluate performance on the AfriMMLU benchmark.
+        Select a language and model to evaluate performance on the AfriMMLU benchmark.
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
-                file_input = gr.File(
-                    label="Upload Test Data (JSON)",
-                    file_types=[".json"]
+                language_input = gr.Dropdown(
+                    choices=list(language_options.keys()),
+                    label="Select Language",
+                    value="swa"
                 )
                 model_input = gr.Dropdown(
                     choices=["deepseek/deepseek-chat"],
@@ -204,14 +266,29 @@ def create_gradio_interface():
                     wrap=True
                 )
         
+        def evaluate_language(language_code, model_name):
+            test_data = load_afrimmlu_data(language_code)
+            if test_data is None:
+                return None, None, None
+            
+            preprocessed_data = preprocess_dataset(test_data)
+            results = evaluate_afrimmlu(preprocessed_data, model_name, language_code)
+            summary_df, plot = create_visualization(results)
+            detailed_df = pd.DataFrame(results["detailed_results"])
+            
+            return summary_df, plot, detailed_df
+        
         evaluate_btn.click(
-            fn=evaluate_and_display,
-            inputs=[file_input, model_input],
+            fn=evaluate_language,
+            inputs=[language_input, model_input],
             outputs=[summary_table, summary_plot, detailed_results]
         )
     
     return demo
 
 if __name__ == "__main__":
+    os.environ['DEEPSEEK_API_KEY']
+    os.environ['HF_TOKEN']
+
     demo = create_gradio_interface()
-    demo.launch(share=True)
+    demo.launch(share=True)