trained with squad dataset

Author: Ubuntu

src/korea_travel_guide/inference.py

@@ -51,15 +51,16 @@ def main():
     inf_args = parse_args()
 
     # set device
-    device = 0 if torch.cuda.is_available() else -1
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     logger.info(f"Using device: {device}")
 
     # load tokenizer + model
     tok = AutoTokenizer.from_pretrained("facebook/bart-base")
     base_model = build_base_model()
-    if training_args.use_sdpa_attention:
+    if inf_args.use_sdpa_attention:
         base_model.config.attn_implementation = "sdpa"
     model = load_peft_model_for_inference(base_model)
+    model.to(device)
 
     # tokenize & format depending on mode
     if inf_args.mode == "test":
@@ -102,7 +103,8 @@ def main():
             return_tensors="pt",
             padding=True,
             truncation=True,
-        ).to(model.device)
+        )
+        enc = {k: v.to(device) for k, v in enc.items()}
 
         # fast batched generate
         out = model.generate(

src/korea_travel_guide/train.py

@@ -34,7 +34,9 @@ class DataArgs:
     train_sample: bool = field(
         default=False, metadata={"help": "Use the mini CSV for smoke tests if True."}
     )
-
+    use_squad: bool = field(
+        default=False, metadata={"help": "If True, ignore CSVs and load SQuAD instead."}
+    )
 
 # training & LoRA extras — extend HF’s own Seq2SeqTrainingArguments
 @dataclass
@@ -126,19 +128,42 @@ def main() -> None:
     logger.info(f"Set seed: {training_args.seed}")
 
     # ---------- Data Preprocessing ----------
-    # load and tokenize dataset
-    # load CSVs
-    data_files = {
-        "train": str(
-            data_args.train_sample_file
-            if data_args.train_sample
-            else data_args.train_file
-        ),
-        "validation": str(data_args.validation_file),
-        "test": str(data_args.test_file),
-    }
-
-    ds = load_dataset("csv", data_files=data_files)
+    # load either CSVs or SQuAD for a quick pipeline sanity check
+    if data_args.use_squad:
+        # 1) pull down SQuAD  
+        raw = load_dataset("squad")
+        # 2) map to simple Q/A pairs (first answer only)
+        def to_qa(ex):
+            return {
+                "question": ex["question"],
+                "answer": ex["answers"]["text"][0]
+            }
+        ds = raw.map(to_qa, remove_columns=raw["train"].column_names)
+    else:
+        # load from your processed CSVs
+        data_files = {
+            "train": str(
+                data_args.train_sample_file
+                if data_args.train_sample
+                else data_args.train_file
+            ),
+            "validation": str(data_args.validation_file),
+            "test": str(data_args.test_file),
+        }
+        ds = load_dataset("csv", data_files=data_files)
+    # # load and tokenize dataset
+    # # load CSVs
+    # data_files = {
+    #     "train": str(
+    #         data_args.train_sample_file
+    #         if data_args.train_sample
+    #         else data_args.train_file
+    #     ),
+    #     "validation": str(data_args.validation_file),
+    #     "test": str(data_args.test_file),
+    # }
+
+    # ds = load_dataset("csv", data_files=data_files)
     ds_tok, tok = tokenize_and_format(ds)
 
     # initialize base model and LoRA