Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -8,7 +8,9 @@ model_id = "Detsutut/Igea-350M-v0.0.1"
|
|
| 8 |
pipeline = transformers.pipeline(
|
| 9 |
"text-generation",
|
| 10 |
model=model_id,
|
| 11 |
-
model_kwargs={"torch_dtype": torch.bfloat16}
|
|
|
|
|
|
|
| 12 |
)
|
| 13 |
|
| 14 |
# Define the function to generate text
|
|
|
|
| 8 |
pipeline = transformers.pipeline(
|
| 9 |
"text-generation",
|
| 10 |
model=model_id,
|
| 11 |
+
model_kwargs={"torch_dtype": torch.bfloat16},
|
| 12 |
+
load_in_8bit=True,
|
| 13 |
+
device_map="auto"
|
| 14 |
)
|
| 15 |
|
| 16 |
# Define the function to generate text
|