JustKiddo commited on
Commit
2650814
1 Parent(s): c99097e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -4,13 +4,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
4
  from huggingface_hub import login
5
  import torch
6
 
 
 
7
  class VietnameseChatbot:
8
- def __init__(self, model_name="TheBloke/Llama-2-7B-Chat-GGML"):
9
  """
10
  Initialize the Vietnamese chatbot with a pre-trained model
11
  """
12
- #self.tokenizer = AutoTokenizer.from_pretrained(model_name)
13
- self.model = AutoModel.from_pretrained("TheBloke/Llama-2-7B-Chat-GGML")
14
 
15
  # Use GPU if available
16
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
4
  from huggingface_hub import login
5
  import torch
6
 
7
+ login(token=os.getenv('HF_TOKEN'))
8
+
9
  class VietnameseChatbot:
10
+ def __init__(self, model_name="meta-llama/Llama-2-7b-chat-hf"):
11
  """
12
  Initialize the Vietnamese chatbot with a pre-trained model
13
  """
14
+ self.tokenizer = AutoTokenizer.from_pretrained(model_name)
15
+ self.model = AutoModel.from_pretrained(model_name)
16
 
17
  # Use GPU if available
18
  self.device = "cuda" if torch.cuda.is_available() else "cpu"