from transformers import AutoTokenizer
from hanuman_loader import HanumanModel


def run(prompt: str = "สวัสดีครับ ช่วยสรุปประเทศไทยสั้น ๆ"):
    repo_id = "ZombitX64/GPT4All-Model"
    tokenizer = AutoTokenizer.from_pretrained(repo_id)
    model_wrapper = HanumanModel.from_pretrained(repo_id, map_location='cpu')
    model = model_wrapper.model
    inputs = tokenizer(prompt, return_tensors='pt')
    out = model.generate(inputs['input_ids'], max_new_tokens=50, temperature=1.2, top_k=50, top_p=0.95)
    print(tokenizer.decode(out[0], skip_special_tokens=True))


if __name__ == '__main__':
    run()