Rajiv714/Axion-0.5B-LiteRT

Static 8 Bit Quantization

Python script to test model

import os
import litert_lm
from huggingface_hub import hf_hub_download

# Set this to either a Hugging Face repo ID or a local .litertlm model path.
MODEL_PATH_OR_REPO = "Rajiv714/Axion-0.5B-LiteRT"
MODEL_FILENAME = "model.litertlm"
MAX_NUM_TOKENS = 1500
TEMPERATURE = 0.3
TOP_K = 40
SYSTEM_PROMPT = "You are a helpful medical AI assistant."
USER_PROMPT = "I am having a fever for the past 10 days. What should I do?"

def get_model_path(model_path_or_repo):
    if os.path.isfile(model_path_or_repo):
        return model_path_or_repo
    print(f"Downloading {MODEL_FILENAME} from {model_path_or_repo}...")

    return hf_hub_download(
        repo_id=model_path_or_repo,
        filename=MODEL_FILENAME,
    )

def main():
    model_path = get_model_path(MODEL_PATH_OR_REPO)
    print(f"Loading model from {model_path} on CPU...\n")

    with litert_lm.Engine(
        model_path, backend=litert_lm.Backend.CPU(), max_num_tokens=MAX_NUM_TOKENS
    ) as engine:
        sampler = litert_lm.SamplerConfig(temperature=TEMPERATURE, top_k=TOP_K)
        system_message = litert_lm.Message.system(SYSTEM_PROMPT)
        with engine.create_conversation(
            messages=[system_message], sampler_config=sampler
        ) as conversation:

            print("--- MODEL RESPONSE ---")
            stream = conversation.send_message_async(
                litert_lm.Message.user(USER_PROMPT)
            )
            for chunk in stream:
                for item in chunk.get("content", []):
                    if item.get("type") == "text":
                        print(item["text"], end="", flush=True)

if __name__ == "__main__":
    main()