Update app.py

app.py

@@ -5,15 +5,21 @@ import torch
 from peft import PeftConfig, PeftModel
 from transformers import LlamaForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-# config = PeftConfig.from_pretrained("GGmorello/FLAMES")
+config = PeftConfig.from_pretrained("GGmorello/FLAMES-20k")
 model = LlamaForCausalLM.from_pretrained(
-    "GGmorello/FLAMES-100k",
+    config.base_model_name_or_path,
+    quantization_config=BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+    ),
 )
-# model = PeftModel.from_pretrained(model, "GGmorello/FLAMES-100k")
+model = PeftModel.from_pretrained(model, "GGmorello/FLAMES-20k")
 
 
 MAX_SEQ_LEN = 4096
-tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-hf")
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 model.config.pad_token = tokenizer.pad_token = tokenizer.unk_token
 
 
@@ -32,4 +38,4 @@ def predict(text):
 
 
 demo = gr.Interface(fn=predict, inputs=gr.Text(), outputs=gr.Text())
-demo.launch()
+demo.launch()