f

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
-import torch
 from datasets import Dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 import pandas as pd
 from huggingface_hub import login
+import torch
 
 def train_model(file, hf_token):
     try:
@@ -16,20 +16,27 @@ def train_model(file, hf_token):
         df = pd.read_csv(file.name)
         dataset = Dataset.from_pandas(df)
         
-        # Model setup
+        # Model setup - force CPU
         model_name = "facebook/opt-125m"
+        device_map = "cpu"  # Force CPU usage
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            device_map=device_map,
+            torch_dtype=torch.float32  # Use float32 for CPU
+        )
         
         # Training configuration
         training_args = TrainingArguments(
             output_dir="./results",
             num_train_epochs=3,
-            per_device_train_batch_size=2,
+            per_device_train_batch_size=1,  # Reduced for CPU
             learning_rate=3e-5,
             save_strategy="epoch",
             push_to_hub=True,
-            hub_token=hf_token
+            hub_token=hf_token,
+            no_cuda=True,  # Force CPU usage
+            report_to="none"  # Disable wandb logging
         )
         
         # Initialize trainer
@@ -43,6 +50,9 @@ def train_model(file, hf_token):
         # Run training
         trainer.train()
         
+        # Push to hub
+        model.push_to_hub(f"cheberle/product-classifier-{pd.Timestamp.now().strftime('%Y%m%d')}")
+        
         return "Training completed successfully!"
         
     except Exception as e:
@@ -57,8 +67,8 @@ demo = gr.Interface(
     ],
     outputs="text",
     title="Product Classifier Training",
-    description="Upload your CSV data to train a product classifier model."
+    description="Upload your CSV data to train a product classifier model on CPU."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=False)