Use transformers to translate English to French

{ pkgs ? import <nixpkgs> {} }:

let
  envDir = "$(pwd)/venv";
in

pkgs.mkShell {
  nativeBuildInputs = [
    pkgs.python310
    pkgs.python310Packages.torch
    pkgs.python310Packages.transformers
    pkgs.python310Packages.sentencepiece
    pkgs.python310Packages.sacremoses
  ];
  shellHook = ''
    export LD_LIBRARY_PATH=${envDir}/lib
    export VIRTUAL_ENV_DISABLE_PROMPT=true

    virtualenv `basename ${envDir}`

    export PIP_PREFIX=${envDir}
    export PYTHONUSERBASE=${envDir}
    export PYTHON_SITE_PACKAGES=$PIP_PREFIX/${pkgs.python310.sitePackages}
    export PYTHONPATH="$PYTHON_SITE_PACKAGES:$PYTHONPATH"
    export PATH="$PIP_PREFIX/bin:${pkgs.ruff}/bin:$PATH"
    unset SOURCE_DATE_EPOCH

    source ${envDir}/bin/activate
  '';
}

import torch
from transformers import MarianMTModel, MarianTokenizer

# Load the pre-trained model and tokenizer
model_name = "Helsinki-NLP/opus-mt-en-fr"
model = MarianMTModel.from_pretrained(model_name)
tokenizer = MarianTokenizer.from_pretrained(model_name)

# Define the input text
input_text = "The quick brown fox jumps over the lazy dog."

# Tokenize the input text
input_ids = tokenizer(input_text, return_tensors="pt").input_ids

# Generate the output text
input_ids = tokenizer(input_text, return_tensors="pt").input_ids
generated_ids = model.generate(input_ids, max_new_tokens=100)

# Decode the generated text
output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

# Print the output text
print(output_text)