Ashishkr commited on
Commit
5d42eda
1 Parent(s): dcb1eb3

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +5 -2
model.py CHANGED
@@ -3,6 +3,9 @@ from typing import Iterator
3
 
4
  import torch
5
  from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 
 
6
 
7
  model_id = 'Ashishkr/llama2_medical_consultation'
8
 
@@ -14,9 +17,9 @@ import torch
14
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
 
16
  config = PeftConfig.from_pretrained("Ashishkr/llama2_medical_consultation")
17
- model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
18
  model = PeftModel.from_pretrained(model, "Ashishkr/llama2_medical_consultation").to(device)
19
- tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
20
 
21
 
22
  def get_prompt(message: str, chat_history: list[tuple[str, str]],
 
3
 
4
  import torch
5
  from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
6
+ import os
7
+ os.environ.getattribute("HF_API_TOKEN")
8
+
9
 
10
  model_id = 'Ashishkr/llama2_medical_consultation'
11
 
 
17
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
18
 
19
  config = PeftConfig.from_pretrained("Ashishkr/llama2_medical_consultation")
20
+ model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", use_auth_token = HF_API_TOKEN)
21
  model = PeftModel.from_pretrained(model, "Ashishkr/llama2_medical_consultation").to(device)
22
+ tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_auth_token = HF_API_TOKEN)
23
 
24
 
25
  def get_prompt(message: str, chat_history: list[tuple[str, str]],