import transformers

model_name = "TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ"
tokenizer_name = "TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ"

model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = transformers.GPT2Tokenizer.from_pretrained(tokenizer_name)

while True:
    user_input = input("You: ")
    if user_input.lower() == "exit":
        break
    encoded_dict = tokenizer.encode_plus(
        user_input,
        max_length=1024,
        pad_to_max_length=True,
        return_attention_mask=True,
        return_tensors="pt",
    )
    output = model.generate(**encoded_dict)
    print("Wizard: ", output[0])