Update hf migrate

Couchbase-Ecosystem · Jan 27, 2025 · d90f014 · d90f014
1 parent 6a073bf
commit d90f014
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 4 deletions.
diff --git a/src/pages/Tools/HuggingFaceMigrate/huggingFaceMigrate.ts b/src/pages/Tools/HuggingFaceMigrate/huggingFaceMigrate.ts
@@ -89,7 +89,8 @@ export const huggingFaceMigrate = async () => {
         case "vscode-couchbase.tools.huggingFaceMigrate.listConfigs":
           try {
             const configs = await HuggingFaceToCouchbase.listConfigs(
-              message.repositoryPath
+              message.repositoryPath,
+              message.trustRemoteCode
             );
             currentPanel.webview.postMessage({
               command:
@@ -109,6 +110,7 @@ export const huggingFaceMigrate = async () => {
           try {
             const splits = await HuggingFaceToCouchbase.listSplits(
               message.repositoryPath,
+              message.trustRemoteCode,
               message.config
             );
             currentPanel.webview.postMessage({
@@ -129,6 +131,7 @@ export const huggingFaceMigrate = async () => {
           try {
             const fields = await HuggingFaceToCouchbase.listFields(
               message.repositoryPath,
+              message.trustRemoteCode,
               message.config,
               message.split
             );
@@ -171,6 +174,7 @@ export const huggingFaceMigrate = async () => {
             const formData = message.data;
               await HuggingFaceToCouchbase.export(
                 formData.repoLink,
+                formData.trustRemoteCode,
                 formData.filePaths,
                 formData.config,
                 formData.split,

diff --git a/src/tools/HFMigrate.ts b/src/tools/HFMigrate.ts
@@ -19,7 +19,8 @@ export class HuggingFaceToCouchbase {
         });
     }
     static async listConfigs(
-        repositoryPath: string
+        repositoryPath: string,
+        trustRemoteCode: string
     ): Promise<string | undefined> {
         const connection = getActiveConnection();
         if (!connection) {
@@ -40,6 +41,10 @@ export class HuggingFaceToCouchbase {
             cmd.push("--path");
             cmd.push(repositoryPath);
             cmd.push('--json-output');
+            if (trustRemoteCode) {
+                cmd.push("--trust-remote-code");
+            }
+
 
             const command = cmd.join(" ");
             const result = await this.runCommand(command); // Run the command and capture the output
@@ -53,6 +58,7 @@ export class HuggingFaceToCouchbase {
 
     static async listSplits(
         repositoryPath: string,
+        trustRemoteCode: string,
         config: string
     ): Promise<string | undefined> {
         const connection = getActiveConnection();
@@ -75,7 +81,10 @@ export class HuggingFaceToCouchbase {
             cmd.push(repositoryPath);
             cmd.push('--name');
             cmd.push(config);
-            cmd.push('--json-output')
+            cmd.push('--json-output');
+            if (trustRemoteCode) {
+                cmd.push("--trust-remote-code");
+            }
 
             const command = cmd.join(" ");
             const result = await this.runCommand(command);
@@ -89,6 +98,7 @@ export class HuggingFaceToCouchbase {
 
     static async listFields(
         repositoryPath: string,
+        trustRemoteCode: string,
         config: string,
         split: string
     ): Promise<string | undefined> {
@@ -115,6 +125,9 @@ export class HuggingFaceToCouchbase {
             cmd.push('--split');
             cmd.push(split);
             cmd.push('--json-output');
+            if (trustRemoteCode) {
+                cmd.push("--trust-remote-code");
+            }
             const command = cmd.join(" ");
             const result = await this.runCommand(command); // Run the command and capture the output
             return result; // Return the output
@@ -127,6 +140,7 @@ export class HuggingFaceToCouchbase {
 
     static async export(
         repoLink: string,
+        trustRemoteCode: string,
         filePaths: string,
         config: string,
         split: string,
@@ -171,6 +185,9 @@ export class HuggingFaceToCouchbase {
                 cmd.push("--split");
                 cmd.push(split);
             }
+            if (trustRemoteCode) {
+                cmd.push("--trust-remote-code");
+            }
             cmd.push("--cb-url");
             cmd.push(connection.url);
             cmd.push("--cb-username");

diff --git a/src/webViews/tools/huggingFaceImporter.webview.ts b/src/webViews/tools/huggingFaceImporter.webview.ts
@@ -314,7 +314,7 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
             <br>
             <div style="display: flex; align-items: center;">
                 <div class="radio-group">
-                    <input type="radio" id="useRepo" name="dataMethod" value="repo">
+                    <input type="radio" id="useRepo" name="dataMethod" value="repo" checked>
                     <label for="useRepo" class="form-label-align">Repo</label>
                 </div>
                 <div class="radio-group">
@@ -335,6 +335,14 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
         <div id="pathInputContainer" style="display:none;">
             <label for="filePaths">File Paths (comma-separated):</label>
             <input type="text" id="filePaths" name="filePaths" placeholder="e.g., /path/to/file1,/path/to/file2">
+        </div>
+        <br>
+            <div class="checkbox-row">
+                <div class="checkbox-container tooltip">
+                    <input type="checkbox" id="trust-remote-code" name="trust-remote-code" checked>
+                    <label for="trust-remote-code">Trust Remote Code</label>
+                    <span class="tooltiptext">Check to enable Trust remote code flag.</span>
+                </div>
         </div>
             <div class="validation-error" id="validation-error-connect"></div>
             <div id="configContainer" style="display:none;">
@@ -451,6 +459,8 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
         event.preventDefault();
         document.getElementById("validation-error").innerHTML = "";
         document.getElementById("validation-error-connect").innerHTML = "";
+        const trustRemoteCodeElement = document.getElementById("trust-remote-code");
+        const trustRemoteCode = trustRemoteCodeElement.checked;
 
         // Ensure the Repo option is selected
         if (!$('#useRepo').is(':checked')) {
@@ -480,6 +490,7 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
         vscode.postMessage({
             command: "vscode-couchbase.tools.huggingFaceMigrate.listConfigs",
             repositoryPath: repoLink,
+            trustRemoteCode: trustRemoteCode,
         });
     }
 
@@ -653,6 +664,8 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
         $('#fields').html(''); // Clear existing options
 
         const repoLink = document.getElementById("repoLink").value;
+        const trustRemoteCodeElement = document.getElementById("trust-remote-code");
+        const trustRemoteCode = trustRemoteCodeElement.checked;
 
         // Disable splits until data is loaded
         $('#splits').prop('disabled', true);
@@ -661,6 +674,7 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
         vscode.postMessage({
             command: "vscode-couchbase.tools.huggingFaceMigrate.listSplits",
             repositoryPath: repoLink,
+            trustRemoteCode: trustRemoteCode,
             config: selectedConfig,
         });
     }
@@ -678,12 +692,15 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
         $('#fields').html(''); // Clear existing options
 
         const repoLink = document.getElementById("repoLink").value;
+        const trustRemoteCodeElement = document.getElementById("trust-remote-code")
+        const trustRemoteCode = trustRemoteCodeElement.checked;
         const selectedConfig = document.getElementById("configs").value;
 
         // Send a message to fetch fields for the selected split
         vscode.postMessage({
             command: "vscode-couchbase.tools.huggingFaceMigrate.listFields",
             repositoryPath: repoLink,
+            trustRemoteCode: trustRemoteCode,
             config: selectedConfig,
             split: selectedSplit,
         });
@@ -734,6 +751,8 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
         document.getElementById("validation-error").innerHTML = "";
 
         const repoLink = document.getElementById("repoLink").value;
+        const trustRemoteCodeCheckBox = document.getElementById("trust-remote-code");
+        const trustRemoteCode = trustRemoteCodeCheckBox.checked;
         const bucket = document.getElementById("bucket").value;
         const scope = document.getElementById("cbScope").value;
         const collection = document.getElementById("cbCollection").value;
@@ -742,6 +761,7 @@ export const huggingFaceMigrateWebView = async (buckets: string[]): Promise<stri
 const dataMethod = $('input[name="dataMethod"]:checked').val();
     let formData = {
         repoLink,
+        trustRemoteCode,
         bucket,
         scope,
         collection,