Update app.py

app.py

@@ -7,8 +7,10 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 # from google.colab import userdata
 import os
 
-#model_id = "somosnlp/gemma-FULL-RAC-Colombia_v2"
-model_id = "somosnlp/RecetasDeLaAbuela_mistral-7b-instruct-v0.2-bnb-4bit"
+#sft_model = "somosnlp/gemma-FULL-RAC-Colombia_v2"
+sft_model = "somosnlp/RecetasDeLaAbuela_mistral-7b-instruct-v0.2-bnb-4bit"
+base_model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -23,14 +25,16 @@ max_seq_length=400
 #     attn_implementation=None
 attn_implementation=None
 
-tokenizer = AutoTokenizer.from_pretrained(model_id,
-                                          max_length = max_seq_length)
-model = AutoModelForCausalLM.from_pretrained(model_id,
-                                            #  quantization_config=bnb_config,
-                                             device_map = {"":0},
-                                              attn_implementation = attn_implementation, # A100 o H100
-                                             ).eval()
+#base_model = AutoModelForCausalLM.from_pretrained(model_name,return_dict=True,torch_dtype=torch.float16,)
+#base_model = AutoModelForCausalLM.from_pretrained(model_name,return_dict=True,device_map="auto", torch_dtype=torch.float16,)
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name, return_dict=True, device_map = {"":0}, attn_implementation = attn_implementation, # A100 o H100).eval()
 
+tokenizer = AutoTokenizer.from_pretrained(sft_model, max_length = max_seq_length)
+ft_model = PeftModel.from_pretrained(base_model, sft_model)
+model = ft_model.merge_and_unload()
+model.save_pretrained(".")
+model.to('cuda')
+tokenizer.save_pretrained(".")
 
 
 class ListOfTokensStoppingCriteria(StoppingCriteria):