Upload README.md with huggingface_hub

README.md

@@ -48,14 +48,23 @@ Trained on the [lucasmccabe/logiqa](https://huggingface.co/datasets/lucasmccabe/
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-model = AutoModelForCausalLM.from_pretrained("rudransh/circuit", torch_dtype=torch.float16)
-tokenizer = AutoTokenizer.from_pretrained("rudransh/circuit")
-
-prompt = "If all squares are rectangles, and all rectangles have four sides, what can we conclude about squares?"
-
-inputs = tokenizer(prompt, return_tensors="pt")
+model = AutoModelForCausalLM.from_pretrained(
+    "rudranshjoshi/circuit",
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
+tokenizer = AutoTokenizer.from_pretrained(
+    "rudranshjoshi/circuit",
+    trust_remote_code=True
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+prompt = "Your prompt here"
+inputs = tokenizer(prompt, return_tensors="pt").to(device)
 outputs = model.generate(**inputs, max_new_tokens=150)
-print(tokenizer.decode(outputs[0], skip_special_tokens=True)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 
 ```