- app.py +8 -8
- requirements.txt +1 -2
app.py
CHANGED
@@ -3,10 +3,10 @@ import pandas as pd
|
|
3 |
from datasets import Dataset
|
4 |
from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForCausalLM
|
5 |
import torch
|
6 |
-
import
|
7 |
|
8 |
-
#
|
9 |
-
|
10 |
|
11 |
def train_model(file, hf_token):
|
12 |
try:
|
@@ -19,10 +19,9 @@ def train_model(file, hf_token):
|
|
19 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
20 |
model = AutoModelForCausalLM.from_pretrained(
|
21 |
model_name,
|
22 |
-
|
23 |
-
|
24 |
)
|
25 |
-
model = model.to('cpu') # Explicitly move to CPU
|
26 |
|
27 |
# Basic dataset creation
|
28 |
dataset = Dataset.from_pandas(df)
|
@@ -32,7 +31,8 @@ def train_model(file, hf_token):
|
|
32 |
per_device_train_batch_size=1,
|
33 |
num_train_epochs=1,
|
34 |
no_cuda=True,
|
35 |
-
local_rank=-1
|
|
|
36 |
)
|
37 |
|
38 |
trainer = Trainer(
|
@@ -58,4 +58,4 @@ demo = gr.Interface(
|
|
58 |
)
|
59 |
|
60 |
if __name__ == "__main__":
|
61 |
-
demo.launch(debug=True)
|
|
|
3 |
from datasets import Dataset
|
4 |
from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForCausalLM
|
5 |
import torch
|
6 |
+
import os
|
7 |
|
8 |
+
# Force CPU
|
9 |
+
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
10 |
|
11 |
def train_model(file, hf_token):
|
12 |
try:
|
|
|
19 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
20 |
model = AutoModelForCausalLM.from_pretrained(
|
21 |
model_name,
|
22 |
+
low_cpu_mem_usage=True,
|
23 |
+
torch_dtype=torch.float32
|
24 |
)
|
|
|
25 |
|
26 |
# Basic dataset creation
|
27 |
dataset = Dataset.from_pandas(df)
|
|
|
31 |
per_device_train_batch_size=1,
|
32 |
num_train_epochs=1,
|
33 |
no_cuda=True,
|
34 |
+
local_rank=-1,
|
35 |
+
use_cpu=True
|
36 |
)
|
37 |
|
38 |
trainer = Trainer(
|
|
|
58 |
)
|
59 |
|
60 |
if __name__ == "__main__":
|
61 |
+
demo.launch(debug=True)
|
requirements.txt
CHANGED
@@ -3,5 +3,4 @@ transformers==4.37.2
|
|
3 |
torch==2.1.2
|
4 |
datasets==2.16.1
|
5 |
pandas==2.2.0
|
6 |
-
huggingface-hub==0.27.0
|
7 |
-
zerogpu
|
|
|
3 |
torch==2.1.2
|
4 |
datasets==2.16.1
|
5 |
pandas==2.2.0
|
6 |
+
huggingface-hub==0.27.0
|
|