JustKiddo commited on
Commit
138c789
1 Parent(s): c25ba20

Update app.py

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