The MedIT One model is an early checkpoint in the development of the One series, evaluated after 9 billion tokens of training.
It is designed for natural language generation tasks and is implemented with a focus on high performance on causal language modeling.
This checkpoint contains 140 million parameters and is built using PyTorch with support for bfloat16 precision, making it suitable for GPU-accelerated inference.
Intended Use
Primary Applications: Natural language generation, research experiments, and prompt completion tasks.
Research: This model checkpoint is provided as an early checkpoint and can be used for studying model behaviors, especially regarding repetitive generation.
Prototyping: Developers and researchers can use this checkpoint to explore early results and understand the evolution of the Medit series.
Caution: As an early checkpoint, the model tends to exhibit repetitive generation. Users should set the repetition penalty (recommended value: 1.2) during inference to mitigate this behavior.
Installation
bash
1# From source (without CUDA acceleration)2git clone https://github.com/MedITSolutionsKurman/medit-one
3cd medit-one
4pip install -e .56# From source with CUDA acceleration7python install_cuda.py
89# For training capabilities only10pip install -e ".[training]"1112# For full installation with all features including CUDA acceleration13pip install -e ".[full]"
How to Use
After installing the medit-one package from the repository, the model can be loaded and run with the following code snippet:
python
1import sys
2import os
3import warnings
45import torch
6from tqdm import tqdm
7import numpy as np
8from transformers import AutoTokenizer, TextStreamer
910from one.modeling_one import OneForCausalLM
1112# Set the model checkpoint path13path ='meditsolutions/medit-one-140M-9B-tokens-checkpoint'1415# Load the tokenizer and model16tokenizer = AutoTokenizer.from_pretrained(path)17model = OneForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16)1819device ='cuda'20model.to(device)2122text ='The role of artificial intelligence'2324# Tokenize input text25tokens = tokenizer(text, return_tensors='pt')26tokens.to(device)2728from time import time
2930start = time()3132# Inference with recommended repetition penalty33with torch.autocast(device_type=device, dtype=torch.bfloat16):34with torch.no_grad():35 model.eval()36 output = model.generate(37**tokens,38 max_new_tokens=1024,39 streamer=TextStreamer(tokenizer),40 do_sample=None,41 temperature=None,42 repetition_penalty=1.2,43 use_cache=True,44 output_attentions=False,45 eos_token_id=model.config.eos_token_id if model.config.eos_token_id isnotNoneelse tokenizer.eos_token_id
46)4748end = time()49tokens_per_sec =len(output[0])/(end - start)50print(f'Time taken: {end - start} seconds, tokens per s: {tokens_per_sec}')
Note: When using this checkpoint, it is essential to apply a repetition penalty of 1.2 to help control the model’s tendency toward repetitive text generation.
Model Details
Parameters: 140M (early checkpoint)
Training Tokens: Evaluated after 9B tokens
Precision: Supports bfloat16 for accelerated computation on compatible hardware
Architecture: Causal language model implemented in PyTorch, part of the MedIT One series
Limitations & Considerations
Repetition: This early checkpoint is known to produce repetitive outputs. Adjusting the repetition penalty (recommended: 1.2) is necessary to reduce this effect.
Early Checkpoint Status: As a checkpoint from an early stage of training, performance and fluency might be lower compared to later, more refined checkpoints.
Usage Recommendations: Best suited for research and experimental purposes rather than production deployment without further fine-tuning.
Training Data & Methodology
While detailed documentation on the training dataset and methods is available in the repository, this checkpoint represents an intermediate stage of training after 9B tokens. Users interested in the training process, dataset specifics, and additional checkpoints are encouraged to consult the repository documentation.
Citation
If you use the Medit One model in your research or applications, please cite the repository:
For more details on installation, model training, and updates, please refer to the repository's README and documentation. Contributions and feedback are welcome from the community.