[cleanup] consolidate some prune_heads logic
Author: sshleifer
Pull request: yes

factors out 5 repetitions  of the following logic
```python
def find_pruneable_heads_and_indices(
    heads: List, n_heads: int, head_size: int, already_pruned_heads: set
) -> Tuple[set, "torch.LongTensor"]:
    mask = torch.ones(n_heads, head_size)
    heads = set(heads) - already_pruned_heads  # Convert to set and remove already pruned heads
    for head in heads:
        # Compute how many pruned heads are before the head and move the index accordingly
        head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
        mask[head] = 0
    mask = mask.view(-1).contiguous().eq(1)
    index: torch.LongTensor = torch.arange(len(mask))[mask].long()
    return heads, index

```

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4799?src=pr&el=h1) Report
> Merging [#4799](https://codecov.io/gh/huggingface/transformers/pull/4799?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/56d5d160cdd177ae6e644506535b56e79feccf68&el=desc) will **decrease** coverage by `0.77%`.
> The diff coverage is `91.30%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/4799/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/4799?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #4799      +/-   ##
==========================================
- Coverage   76.15%   75.38%   -0.78%     
==========================================
  Files         128      128              
  Lines       21497    21464      -33     
==========================================
- Hits        16371    16180     -191     
- Misses       5126     5284     +158     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/4799?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_gpt2.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19ncHQyLnB5) | `31.93% <50.00%> (-53.83%)` | :arrow_down: |
| [src/transformers/modeling\_t5.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190NS5weQ==) | `84.26% <50.00%> (+0.99%)` | :arrow_up: |
| [src/transformers/modeling\_albert.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19hbGJlcnQucHk=) | `76.33% <100.00%> (-0.33%)` | :arrow_down: |
| [src/transformers/modeling\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iZXJ0LnB5) | `86.27% <100.00%> (-0.16%)` | :arrow_down: |
| [src/transformers/modeling\_distilbert.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19kaXN0aWxiZXJ0LnB5) | `97.52% <100.00%> (-0.05%)` | :arrow_down: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `80.96% <100.00%> (+0.96%)` | :arrow_up: |
| [src/transformers/modeling\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ191dGlscy5weQ==) | `90.02% <100.00%> (-0.14%)` | :arrow_down: |
| [src/transformers/modeling\_xlm.py](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ194bG0ucHk=) | `89.27% <100.00%> (-0.16%)` | :arrow_down: |
| ... and [2 more](https://codecov.io/gh/huggingface/transformers/pull/4799/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/4799?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4799?src=pr&el=footer). Last update [56d5d16...bfb3251](https://codecov.io/gh/huggingface/transformers/pull/4799?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Nice! LGTM