Update model.py
Browse files
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]],
|