model.py

import torch
from torch import nn
from transformers import AutoTokenizer, T5ForConditionalGeneration

device = "cuda" if torch.cuda.is_available() else "cpu"

def create_flan_T5_model(device=device):
  """Creates a HuggingFace all-MiniLM-L6-v2 model.

  Args:
    device: A torch.device
  Returns:
    A tuple of the model and tokenizer
  """
  tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-small')
  model = T5ForConditionalGeneration.from_pretrained('google/flan-t5-small').to(device)

  return model, tokenizer

# Example usage
model, tokenizer = create_flan_T5_model()