Update fetch_data.py

fetch_data.py

@@ -1,8 +1,16 @@
 import requests
 import pandas as pd
+import json
 
-def fetch_and_update_training_data(train_file):
-    model_name = "TYT"
+# Load configuration file 
+with open('config.json', 'r') as config_file:
+    config = json.load(config_file) 
+
+arg2 = config.get('arg2', 'train_fleet') 
+arg7 = config.get('arg7', 'Model')
+
+def fetch_and_update_training_data(arg2):
+    model_name = arg7
     export_type = "sentence"
     request_url = f"http://20.247.235.135/AiBot/BE/api/PromptEnhancer/ExportPEList?modelName={model_name}&exportType={export_type}"
 
@@ -11,13 +19,13 @@ def fetch_and_update_training_data(train_file):
     if response.status_code == 200:
         try:
             # Save response as CSV
-            with open(train_file + ".csv", "w", encoding="utf-8") as file:
+            with open(arg2 + ".csv", "w", encoding="utf-8") as file:
                 file.write(response.text)
 
             print("Training file updated successfully.")
 
             # Read the updated CSV into a DataFrame
-            return pd.read_csv(train_file + ".csv")
+            return pd.read_csv(arg2 + ".csv")
 
         except Exception as e:
             print(f"Error processing the response: {e}")