Update app.py
Browse files
app.py
CHANGED
@@ -4,9 +4,6 @@ import torch
|
|
4 |
from transformers import pipeline
|
5 |
|
6 |
|
7 |
-
@spaces.GPU
|
8 |
-
def warm_up():
|
9 |
-
pass
|
10 |
|
11 |
# Set up the model
|
12 |
model_id = "meta-llama/Llama-3.2-1B"
|
@@ -14,14 +11,13 @@ pipe = pipeline(
|
|
14 |
"text-generation",
|
15 |
model=model_id,
|
16 |
torch_dtype=torch.bfloat16,
|
17 |
-
device_map="
|
18 |
)
|
19 |
|
20 |
|
21 |
|
22 |
-
|
23 |
def generate_text(prompt):
|
24 |
-
pipe.model.to('cuda')
|
25 |
response = pipe(prompt, max_length=100, num_return_sequences=1)
|
26 |
return response[0]['generated_text']
|
27 |
|
@@ -30,7 +26,7 @@ css = """
|
|
30 |
body {background-color: #f0f8ff; font-family: 'Arial', sans-serif;}
|
31 |
.gradio-container {max-width: 700px; margin: 0 auto; border-radius: 15px; padding: 20px; background-color: white; box-shadow: 0 4px 8px rgba(0,0,0,0.1);}
|
32 |
textarea {border-radius: 10px; padding: 10px; font-size: 16px; border: 2px solid #ddd;}
|
33 |
-
button {background-color: #4caf50; color:
|
34 |
button:hover {background-color: #45a049;}
|
35 |
h1 {color: #333; font-size: 32px; text-align: center;}
|
36 |
footer {text-align: center; padding: 10px;}
|
|
|
4 |
from transformers import pipeline
|
5 |
|
6 |
|
|
|
|
|
|
|
7 |
|
8 |
# Set up the model
|
9 |
model_id = "meta-llama/Llama-3.2-1B"
|
|
|
11 |
"text-generation",
|
12 |
model=model_id,
|
13 |
torch_dtype=torch.bfloat16,
|
14 |
+
device_map="cuda"
|
15 |
)
|
16 |
|
17 |
|
18 |
|
19 |
+
@spaces.GPU
|
20 |
def generate_text(prompt):
|
|
|
21 |
response = pipe(prompt, max_length=100, num_return_sequences=1)
|
22 |
return response[0]['generated_text']
|
23 |
|
|
|
26 |
body {background-color: #f0f8ff; font-family: 'Arial', sans-serif;}
|
27 |
.gradio-container {max-width: 700px; margin: 0 auto; border-radius: 15px; padding: 20px; background-color: white; box-shadow: 0 4px 8px rgba(0,0,0,0.1);}
|
28 |
textarea {border-radius: 10px; padding: 10px; font-size: 16px; border: 2px solid #ddd;}
|
29 |
+
button {background-color: #4caf50; color: #000000; border-radius: 5px; padding: 10px 20px; font-size: 16px; border: none; cursor: pointer;}
|
30 |
button:hover {background-color: #45a049;}
|
31 |
h1 {color: #333; font-size: 32px; text-align: center;}
|
32 |
footer {text-align: center; padding: 10px;}
|