Update app.py

app.py

@@ -26,7 +26,7 @@ def process(action, base_model_name, ft_model_name, dataset_name, system_prompt,
         result = prompt_model(ft_model_name, system_prompt, user_prompt, sql_schema)
     return result
 
-def fine_tune_model(model_name, dataset_name):
+def fine_tune_model(base_model_name, dataset_name):
     # Load dataset
     dataset = load_dataset(dataset_name)
 
@@ -35,7 +35,7 @@ def fine_tune_model(model_name, dataset_name):
     print("###")
     
     # Load model
-    model, tokenizer = load_model(model_name)
+    model, tokenizer = load_model(base_model_name)
 
     print("### Model")
     print(model)
@@ -128,7 +128,7 @@ def prompt_model(model_name, system_prompt, user_prompt, sql_schema):
 
 def load_model(model_name):
     model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(model)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
     
     return model, tokenizer