Update app.py

app.py

@@ -12,9 +12,10 @@ from pydantic import BaseModel
 hf_login(token=os.getenv("HF_TOKEN"))
 
 model_name = "meta-llama/Llama-3.2-1B-Instruct"
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 model = AutoModelForCausalLM.from_pretrained(
-    model_name, torch_dtype=torch.float32, device_map="auto"
+    model_name, torch_dtype=torch.float32, device_map=device
 )
 
 class Person(BaseModel):
@@ -88,7 +89,7 @@ def summarize(text):
         {"role": "user", "content": text},
     ]
     
-    model_inputs = tokenizer([text], return_tensors="pt").to('cuda')
+    model_inputs = tokenizer([text], return_tensors="pt").to(device)
 
     generated_ids = model.generate(
         input_ids=model_inputs["input_ids"],
@@ -119,13 +120,6 @@ with gr.Blocks() as demo:
             max_lines=15, 
             placeholder="Paste your clinical report here...",
         )
-        output_text = gr.Textbox(
-            label="Summary", 
-            autoscroll=False,
-            lines=15, 
-            max_lines=15, 
-            show_copy_button=True,
-        )
 
     with gr.Row():
         summarize_btn = gr.Button("Summarize")
@@ -136,9 +130,18 @@ with gr.Blocks() as demo:
             show_progress=True,
         )
 
+    with gr.Row():
+        output_text = gr.Textbox(
+            label="Summary", 
+            autoscroll=False,
+            lines=15, 
+            max_lines=15, 
+            show_copy_button=True,
+        )
+
     with gr.Row():
         examples = gr.Examples(
-            label="Example – Dialog",
+            label="Examples",
             examples=[
 """A 48-year-old Indian woman was referred to the department of periodontology, with the chief complaint of swollen gums. She felt discomfort with the disfigurement of gums which appeared un esthetic due to its more severity and there was bleeding and difficulty while chewing food.
 Past medical history revealed that she is under medication for hypertension with amlodipine (10 mg/day orally) for the last 2 years and 6 months. The amlodipine dose was increased to 50 mg/day orally and statins were prescribed due to the acute angina attack and hypercholesterolemia before 6 months of the dental visit.