Update pbs_data.py

pbs_data.py

@@ -1,13 +1,11 @@
+import datetime
 import requests
 import csv
 from io import StringIO
-import json
 import time
 from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry
-import sqlite3
-import datetime
-import os
+from datasets import Dataset
 
 class PBSPublicDataAPIClient:
     def __init__(self, subscription_key, base_url='https://data-api.health.gov.au/pbs/api/v3', rate_limit=0.2):
@@ -389,30 +387,16 @@ class PBSPublicDataAPIClient:
 
     def preprocess_data(self, data):
         processed = {
-            'drugs': set(),
-            'brands': set(),
-            'formulations': set(),
-            'indications': set(),
-            'treatment_phases': set(),
-            'hospital_types': set(),
             'combinations': []
         }
         
         for pbs_code, item in data.items():
-            processed['drugs'].add(item['name'])
-            processed['brands'].update(item['brands'])  # Update this line
-            processed['formulations'].add(item['li_form'])
-            processed['hospital_types'].add(item['hospital_type'])
-            
             for restriction in item['restrictions']:
-                processed['indications'].add(restriction['indications'])
-                processed['treatment_phases'].add(restriction['treatment_phase'])
-                
-                for brand in item['brands']:  # Add this loop
+                for brand in item['brands']:
                     processed['combinations'].append({
                         'pbs_code': pbs_code,
                         'drug': item['name'],
-                        'brand': brand,  # Update this line
+                        'brand': brand,
                         'formulation': item['li_form'],
                         'indication': restriction['indications'],
                         'treatment_phase': restriction['treatment_phase'],
@@ -425,76 +409,18 @@ class PBSPublicDataAPIClient:
                         'schedule_month': item['schedule_month']
                     })
         
-        return {k: sorted(v) if isinstance(v, set) else v for k, v in processed.items()}
+        return processed
 
-    def save_data_to_sqlite(self, data, db_path="rheumatology_biologics_data.db"):
+    def save_data_to_hf(self, data, dataset_name="cmcmaster/rheumatology-biologics-dataset"):
         processed_data = self.preprocess_data(data)
-
-        # Remove the existing database file if it exists
-        if os.path.exists(db_path):
-            os.remove(db_path)
-
-        conn = sqlite3.connect(db_path)
-        cursor = conn.cursor()
-
-        # Create tables
-        cursor.execute('''CREATE TABLE IF NOT EXISTS drugs
-                          (id INTEGER PRIMARY KEY, name TEXT UNIQUE)''')
-        cursor.execute('''CREATE TABLE IF NOT EXISTS brands
-                          (id INTEGER PRIMARY KEY, name TEXT UNIQUE)''')
-        cursor.execute('''CREATE TABLE IF NOT EXISTS formulations
-                          (id INTEGER PRIMARY KEY, name TEXT UNIQUE)''')
-        cursor.execute('''CREATE TABLE IF NOT EXISTS indications
-                          (id INTEGER PRIMARY KEY, name TEXT UNIQUE)''')
-        cursor.execute('''CREATE TABLE IF NOT EXISTS treatment_phases
-                          (id INTEGER PRIMARY KEY, name TEXT UNIQUE)''')
-        cursor.execute('''CREATE TABLE IF NOT EXISTS hospital_types
-                          (id INTEGER PRIMARY KEY, name TEXT UNIQUE)''')
-        cursor.execute('''CREATE TABLE IF NOT EXISTS combinations
-                          (id INTEGER PRIMARY KEY, pbs_code TEXT, drug_id INTEGER, brand_id INTEGER,
-                           formulation_id INTEGER, indication_id INTEGER, treatment_phase_id INTEGER,
-                           streamlined_code TEXT, online_application BOOLEAN, authority_method TEXT,
-                           hospital_type_id INTEGER, schedule_code TEXT, schedule_year INTEGER, schedule_month TEXT,
-                           FOREIGN KEY (drug_id) REFERENCES drugs(id),
-                           FOREIGN KEY (brand_id) REFERENCES brands(id),
-                           FOREIGN KEY (formulation_id) REFERENCES formulations(id),
-                           FOREIGN KEY (indication_id) REFERENCES indications(id),
-                           FOREIGN KEY (treatment_phase_id) REFERENCES treatment_phases(id),
-                           FOREIGN KEY (hospital_type_id) REFERENCES hospital_types(id))''')
-
-        # Insert data
-        for table in ['drugs', 'brands', 'formulations', 'indications', 'treatment_phases', 'hospital_types']:
-            cursor.executemany(f"INSERT OR IGNORE INTO {table} (name) VALUES (?)",
-                               [(item,) for item in processed_data[table]])
-
-        # Insert combinations
-        for combo in processed_data['combinations']:
-            cursor.execute('''INSERT INTO combinations
-                              (pbs_code, drug_id, brand_id, formulation_id, indication_id,
-                               treatment_phase_id, streamlined_code, online_application, authority_method,
-                               hospital_type_id, schedule_code, schedule_year, schedule_month)
-                              VALUES (?, 
-                                      (SELECT id FROM drugs WHERE name = ?),
-                                      (SELECT id FROM brands WHERE name = ?),
-                                      (SELECT id FROM formulations WHERE name = ?),
-                                      (SELECT id FROM indications WHERE name = ?),
-                                      (SELECT id FROM treatment_phases WHERE name = ?),
-                                      ?, ?, ?,
-                                      (SELECT id FROM hospital_types WHERE name = ?),
-                                      ?, ?, ?)''',
-                           (combo['pbs_code'], combo['drug'], combo['brand'], combo['formulation'],
-                            combo['indication'], combo['treatment_phase'], combo['streamlined_code'],
-                            combo['online_application'], combo['authority_method'], combo['hospital_type'],
-                            combo['schedule_code'], combo['schedule_year'], combo['schedule_month']))
-
-        # Add last_updated column and insert timestamp
-        cursor.execute('''CREATE TABLE IF NOT EXISTS metadata
-                          (key TEXT PRIMARY KEY, value TEXT)''')
-        cursor.execute('''INSERT OR REPLACE INTO metadata (key, value)
-                          VALUES ('last_updated', ?)''', (datetime.datetime.now().isoformat(),))
-
-        conn.commit()
-        conn.close()
+        
+        # Create a Dataset from the combinations
+        dataset = Dataset.from_list(processed_data['combinations'])
+        
+        # Push the dataset to the Hugging Face Hub
+        dataset.push_to_hub(dataset_name)
+        
+        print(f"Data saved to Hugging Face Hub: {dataset_name}")
 
 def main():
     client = PBSPublicDataAPIClient("2384af7c667342ceb5a736fe29f1dc6b", rate_limit=0.2)
@@ -505,8 +431,8 @@ def main():
         
         print(f"Data fetched for {len(data)} items.")
         
-        client.save_data_to_sqlite(data)
-        print("Data saved to rheumatology_biologics_data.db")
+        client.save_data_to_hf(data)
+        print("Data saved to Hugging Face Hub")
 
     except Exception as e:
         print(f"An error occurred: {str(e)}")