Update app.py
Browse files
app.py
CHANGED
@@ -13,13 +13,13 @@ from pydantic import BaseModel
|
|
13 |
hf_login(token=os.getenv("HF_TOKEN"))
|
14 |
|
15 |
model_name = "meta-llama/Llama-3.2-1B-Instruct"
|
16 |
-
model_name = "gregorlied/Llama-3.2-1B-Instruct-Medical-Report-Summarization"
|
17 |
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
18 |
|
19 |
model = AutoModelForCausalLM.from_pretrained(
|
20 |
model_name,
|
21 |
device_map="auto",
|
22 |
-
torch_dtype=torch.bfloat16,
|
23 |
attn_implementation='eager',
|
24 |
trust_remote_code=True,
|
25 |
)
|
@@ -93,10 +93,10 @@ Please extract relevant clinical information from the report.
|
|
93 |
def summarize(text):
|
94 |
if not text.strip():
|
95 |
return "Please enter some text to summarize."
|
96 |
-
|
97 |
messages = [
|
98 |
-
{"role": "system", "content": prompt},
|
99 |
-
{"role": "user", "content": text},
|
100 |
]
|
101 |
|
102 |
model_inputs = tokenizer([text], return_tensors="pt").to(device)
|
|
|
13 |
hf_login(token=os.getenv("HF_TOKEN"))
|
14 |
|
15 |
model_name = "meta-llama/Llama-3.2-1B-Instruct"
|
16 |
+
model_name = "gregorlied/Llama-3.2-1B-Instruct-Medical-Report-Summarization-FP32"
|
17 |
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
18 |
|
19 |
model = AutoModelForCausalLM.from_pretrained(
|
20 |
model_name,
|
21 |
device_map="auto",
|
22 |
+
# torch_dtype=torch.bfloat16,
|
23 |
attn_implementation='eager',
|
24 |
trust_remote_code=True,
|
25 |
)
|
|
|
93 |
def summarize(text):
|
94 |
if not text.strip():
|
95 |
return "Please enter some text to summarize."
|
96 |
+
|
97 |
messages = [
|
98 |
+
{"role": "system", "content": prompt.strip()},
|
99 |
+
{"role": "user", "content": text.strip()},
|
100 |
]
|
101 |
|
102 |
model_inputs = tokenizer([text], return_tensors="pt").to(device)
|