Update app.py

app.py

@@ -3,7 +3,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig,B
 
 import torch
 
-model_id = "truongghieu/deci-finetuned_BK_Regulation"
+model_id = "truongghieu/deci-finetuned_Prj2"
 
 # Check if a GPU is available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -16,14 +16,16 @@ bnb_config = BitsAndBytesConfig(
 
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 # Load model in this way if use GPU
-# model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, quantization_config=bnb_config)
-model = AutoModelForCausalLM.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True)
+if device == "cuda":
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, quantization_config=bnb_config)
+else:
+    model = AutoModelForCausalLM.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True)
 # Move the model to the GPU if available
 
 generation_config = GenerationConfig(
     penalty_alpha=0.6,
     do_sample=True,
-    top_k=5,
+    top_k=3,
     temperature=0.5,
     repetition_penalty=1.2,
     max_new_tokens=50,