yakine commited on
Commit
f6629ad
·
verified ·
1 Parent(s): 8d08ae6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -27,12 +27,14 @@ text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokeniz
27
 
28
  # Load the Llama model and tokenizer at the beginning
29
  model_name = "meta-llama/Meta-Llama-3.1-8B"
 
30
  model_llama = AutoModelForCausalLM.from_pretrained(
31
  model_name,
32
- torch_dtype=torch.float16, # Use FP16 for reduced memory
 
33
  token=hf_token
34
  )
35
- tokenizer_llama = AutoTokenizer.from_pretrained(model_name, token=hf_token)
36
 
37
 
38
 
 
27
 
28
  # Load the Llama model and tokenizer at the beginning
29
  model_name = "meta-llama/Meta-Llama-3.1-8B"
30
+ tokenizer_llama = AutoTokenizer.from_pretrained(model_name, token=hf_token)
31
  model_llama = AutoModelForCausalLM.from_pretrained(
32
  model_name,
33
+ torch_dtype="auto",
34
+ device_map="auto",
35
  token=hf_token
36
  )
37
+
38
 
39
 
40