Update app.py

app.py

@@ -1,7 +1,5 @@
 import torch
 import gradio as gr
-import os
-import logging
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -10,12 +8,11 @@ from transformers import (
     DataCollatorForLanguageModeling
 )
 from datasets import load_dataset
+import logging
+import os
 
-# Force CPU-only mode
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-os.environ["BITSANDBYTES_NOWELCOME"] = "1"
-
-# Configure logging
+# Configure environment
+os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Force CPU
 logging.basicConfig(level=logging.INFO)
 
 def train():
@@ -26,17 +23,13 @@ def train():
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             device_map="cpu",
-            trust_remote_code=True,
-            load_in_4bit=False  # Disable quantization
+            trust_remote_code=True
         )
 
-        # Add padding token
-        tokenizer.pad_token = tokenizer.eos_token
-
-        # Load sample dataset
+        # Load dataset
         dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
 
-        # Tokenization function
+        # Tokenization
         def tokenize_function(examples):
             return tokenizer(
                 examples["text"],
@@ -52,25 +45,21 @@ def train():
             remove_columns=["text"]
         )
 
-        # Data collator
+        # Training setup
         data_collator = DataCollatorForLanguageModeling(
             tokenizer=tokenizer,
             mlm=False
         )
 
-        # Training arguments
         training_args = TrainingArguments(
             output_dir="./results",
             per_device_train_batch_size=2,
-            per_device_eval_batch_size=2,
-            num_train_epochs=1,  # Reduced for testing
+            num_train_epochs=1,
             logging_dir="./logs",
             fp16=False,
-            bf16=False,
-            use_cpu=True  # Explicit CPU usage
+            report_to="none"
         )
 
-        # Trainer
         trainer = Trainer(
             model=model,
             args=training_args,
@@ -79,11 +68,11 @@ def train():
         )
 
         # Start training
-        logging.info("Starting training...")
+        logging.info("Training started...")
         trainer.train()
         logging.info("Training completed!")
         
-        return "✅ Training successful! Model saved."
+        return "✅ Training successful"
 
     except Exception as e:
         logging.error(f"Error: {str(e)}")