A bare-bones CLI for interacting with the `GPT4All-J@v1.2-jazzy` LLM.

gpt4all_j.py

from fire import Fire
import torch
from transformers import pipeline


def main():
    generate_text = pipeline(
        model="nomic-ai/gpt4all-j",
        revision="v1.2-jazzy",
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto",
        max_length=2000,
    )
    print("Welcome to GPT4All-J!")
    prompt = ""
    while prompt.lower() not in {"q", "quit", "exit"}:
        prompt = input(">>> ")
        res = generate_text(prompt)
        output: str = res[0]["generated_text"]
        content = output[len(prompt) :].strip()
        print(content)


if __name__ == "__main__":
    Fire(main)