import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

peft_model_id = "ntedeschi/reconcile_the_irreconcilable"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=False, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id)

from IPython.display import display, Markdown

def make_inference(topic):
  batch = tokenizer(
      f"### INSTRUCTION\nBelow is a philosophy topic. Please write Hegel's view on the topic, Ayn Rand's view \
      on the topic and a reconciliation of their views. \
      \n\n### Topic:\n{topic}\n \
      \n\n### Hegel:\n \
      \n\n### Ayn Rand:\n \
      \n\n### Reconciliation:\n",
      return_tensors='pt'
  )


  with torch.cuda.amp.autocast():
    output_tokens = model.generate(**batch, max_new_tokens=512)

  display(Markdown((tokenizer.decode(output_tokens[0], skip_special_tokens=True))))

philosophy_topic = "Mind body dualism"
make_inference(philosophy_topic)