fix: add missing imports

aitomatic · Oct 26, 2024 · dd9e994 · dd9e994
1 parent cad014a
commit dd9e994
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/raw_inference.py b/raw_inference.py
@@ -1,3 +1,12 @@
+import argparse
+import logging
+
+import torch
+import yaml
+from peft import LoraConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
+
+
 def load_config(config_file="./configs/inference-config.yaml"):
     """
     Load the configuration from a YAML file.
@@ -109,7 +118,7 @@ def main():
         device_map=config["model"]["device_map"],
     )
 
-    model = PeftModel.from_pretrained(base_model, config["model"]["output_dir"])
+    model = PeftModel1.from_pretrained(base_model, config["model"]["output_dir"])
     model = model.merge_and_unload()
 
     # Perform text generation

diff --git a/training.py b/training.py
@@ -4,12 +4,7 @@
 import yaml
 from datasets import load_dataset
 from peft import LoraConfig
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig,
-    TrainingArguments,
-)
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
 from trl import SFTTrainer