Repository: EleutherAI/lm-evaluation-harness
Issue: Integrating MaxText with lm-evaluation-harness
Root cause: Documentation Deficiency

Hi team,

We are interested in integrating our JAX-based [MaxText repository](https://github.com/AI-Hypercomputer/maxtext) with the lm-evaluation-harness for model evaluation. MaxText serves as a reference implementation for models running on Google Cloud TPUs and GPUs.

Could you please provide guidance on the necessary steps for this integration?

Thank you!

Hi! Thanks for your interest! 

Generally the starting point is to create a wrapper `LM` class that implements the three core methods of `loglikelihood`, `loglikelihood_rolling` and `generate_until` (some details [here](https://github.com/EleutherAI/lm-evaluation-harness/blob/main/docs/model_guide.md)). These methods are expected to receive _all the_ requests and return the corresponding outputs.

Currently the main entrypoint to evaluations is the [evaluate](https://github.com/EleutherAI/lm-evaluation-harness/blob/6b3f3f7e2aece0597a0a4853d364061f93eef5ef/lm_eval/evaluator.py#L405) function, which orchestrates the evaluation process, calls the appropriate model methods and handles result aggregation across (possibly distributed ranks). _However_, the current implementation is quite tightly coupled to PyTorch and HuggingFace conventions, but we're actively working on simplifying and modularizing the pipeline [see #3083] which should make the integration much smoother.

Feel free to drop by the `#lm-thunderdome` channel on [EleutherAI Discord](https://discord.com/channels/729741769192767510/755950983669874798) or ping me there (`baber_`) if you'd like to discuss this further or have any other questions. Would also be happy to create a dedicated channel!


Thank you @baberabb for the info! Given the current tight coupling to PyTorch and Hugging Face, it sounds like running JAX-based models could lead to dependency issues. Do you anticipate simplifying the pipeline/architecture would help resolve this?

1) What's your rough estimate for the timeline on pipeline/architecture simplification?
2) What's your rough estimate for how long it would take users to integrate a model from a different framework with this tool?



wrt to 1. I'm aiming to get this done by August at the latest.  As for 2, I don't expect the overall framework to change dramatically. We've traditionally maintained a clean separation between tasks and models - the key is that models receive strings as input to the appropriate methods (`loglikelihood`, `loglikelihood_rolling`, `generate_until`) and return the expected outputs. Once you have a wrapper that implements these methods for your JAX model, the integration should be relatively straightforward. 

It's very much doable now but the main reason we recommend `evaluate` as an entry point is that it handles the utility work of sending requests to the model (using optional model specific runtime parameters mostly related to chat templating) and calling the appropriate task methods for post-processing. But mainly it's to ensure forward compatibility as we evolve the framework. 

@Helw150 has had a lot of success using this library with his JAX-based models and may be able to provide pointers / guidance.

We don't have anything with MaxText but you can read how we are doing evaluation on TPUs here! https://marin.readthedocs.io/en/latest/explanations/evaluation/

Our setup is a bit janky since we use our custom evaluator for MCQ and VLLM for generation

If you just want a reference for a custom evaluator, see https://github.com/stanford-crfm/levanter/blob/main/src%2Flevanter%2Feval_harness.py