Spaces:
Runtime error
Runtime error
| 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() | |