JustKiddo commited on
Commit
1934f9e
1 Parent(s): 2650814

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -4
app.py CHANGED
@@ -4,15 +4,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
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"
 
4
  from huggingface_hub import login
5
  import torch
6
 
 
 
7
  class VietnameseChatbot:
8
+ def __init__(self, model_name="tamgrnguyen/Gemma-2-2b-it-Vietnamese-Aesthetic"):
9
  """
10
  Initialize the Vietnamese chatbot with a pre-trained model
11
  """
12
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)
13
+ self.model = AutoModelForCausalLM.from_pretrained(model_name)
14
 
15
  # Use GPU if available
16
  self.device = "cuda" if torch.cuda.is_available() else "cpu"