Update process_hf_dataset.py

process_hf_dataset.py

@@ -36,7 +36,7 @@ def rename_variables(code, variable_prefixes=None):
         # Determine variable role based on context (simplified)
         if var in ['expression', 'input']:  # Assume input parameters or initial variables
             role = 'input_variable'
-        elif var in code.split() and 'return' in line for line in code_lines if var in line:  # Returned variables
+        elif var in code.split() and 'return' in line[x] for x in code_lines if var in line[x]:  # Returned variables
             role = 'returned_variable'
         else:  # Default to assigned variables
             role = 'assigned_variable'