First commit

app.py

@@ -1,64 +1,284 @@
+# app.py
+import os
+import json
+import torch
+import pandas as pd
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+from sqlalchemy import create_engine, text
+from transformers import (
+    TrainingArguments,
+    Trainer,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling
+)
+from datasets import Dataset
+from peft import (
+    prepare_model_for_kbit_training,
+    LoraConfig,
+    get_peft_model
 )
+from datetime import datetime
+
+# Constants - Modified for HF Spaces
+MODEL_NAME = "deepseek-ai/DeepSeek-R1"
+OUTPUT_DIR = "/tmp/finetuned_models"  # Using /tmp for HF Spaces
+LOGS_DIR = "/tmp/training_logs"       # Using /tmp for HF Spaces
+
+class TrainingInterface:
+    def __init__(self):
+        self.current_status = "Idle"
+        self.progress = 0
+        self.is_training = False
+        
+    def get_database_url(self):
+        """Get database URL from HF Space secrets"""
+        database_url = os.environ.get('DATABASE_URL')
+        if not database_url:
+            raise Exception("DATABASE_URL not found in environment variables")
+        return database_url
+
+    def fetch_training_data(self, progress=gr.Progress()):
+        """Fetch training data from database"""
+        try:
+            database_url = self.get_database_url()
+            engine = create_engine(database_url)
+            
+            progress(0, desc="Connecting to database...")
+            
+            with engine.connect() as conn:
+                result = conn.execute(text("SELECT COUNT(*) FROM bents"))
+                total_rows = result.scalar()
+                
+                query = text("SELECT chunk_id, text FROM bents")
+                df = pd.read_sql_query(query, conn)
+            
+            progress(0.5, desc="Data fetched successfully")
+            return df
+            
+        except Exception as e:
+            raise gr.Error(f"Database error: {str(e)}")
+
+    def prepare_training_data(self, df, progress=gr.Progress()):
+        """Convert DataFrame into training format"""
+        formatted_data = []
+        try:
+            total_rows = len(df)
+            for idx, row in enumerate(df.iterrows()):
+                progress(idx/total_rows, desc="Preparing training data...")
+                _, row_data = row
+                chunk_id = str(row_data['chunk_id']).strip()
+                text = str(row_data['text']).strip()
+                
+                if chunk_id and text:
+                    formatted_text = f"User: {chunk_id}\nAssistant: {text}"
+                    formatted_data.append({"text": formatted_text})
+            
+            if not formatted_data:
+                raise ValueError("No valid training data found")
+                
+            return formatted_data
+        except Exception as e:
+            raise gr.Error(f"Data preparation error: {str(e)}")
+
+    def stop_training(self):
+        """Stop the training process"""
+        self.is_training = False
+        return "Training stopped by user."
+
+    def train_model(
+        self,
+        learning_rate=2e-4,
+        num_epochs=3,
+        batch_size=4,
+        progress=gr.Progress()
+    ):
+        """Main training function"""
+        try:
+            self.is_training = True
+            
+            # Create directories in /tmp for HF Spaces
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            specific_output_dir = os.path.join(OUTPUT_DIR, f"run_{timestamp}")
+            os.makedirs(specific_output_dir, exist_ok=True)
+            os.makedirs(LOGS_DIR, exist_ok=True)
+
+            # Data preparation
+            progress(0.1, desc="Fetching data...")
+            if not self.is_training:
+                return "Training cancelled."
+                
+            df = self.fetch_training_data()
+            formatted_data = self.prepare_training_data(df)
+            
+            # Model initialization
+            progress(0.2, desc="Loading model...")
+            if not self.is_training:
+                return "Training cancelled."
+                
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                trust_remote_code=True,
+                torch_dtype=torch.float16,
+                load_in_8bit=True,
+                device_map="auto"  # Important for HF Spaces GPU allocation
+            )
+
+            # LoRA configuration
+            progress(0.3, desc="Setting up LoRA...")
+            if not self.is_training:
+                return "Training cancelled."
+                
+            lora_config = LoraConfig(
+                r=16,
+                lora_alpha=32,
+                target_modules=[
+                    "q_proj", "k_proj", "v_proj", "o_proj",
+                    "gate_proj", "up_proj", "down_proj"
+                ],
+                lora_dropout=0.05,
+                bias="none",
+                task_type="CAUSAL_LM"
+            )
+
+            model = prepare_model_for_kbit_training(model)
+            model = get_peft_model(model, lora_config)
+
+            # Training setup
+            progress(0.4, desc="Configuring training...")
+            if not self.is_training:
+                return "Training cancelled."
+                
+            training_args = TrainingArguments(
+                output_dir=specific_output_dir,
+                num_train_epochs=num_epochs,
+                per_device_train_batch_size=batch_size,
+                learning_rate=learning_rate,
+                fp16=True,
+                gradient_accumulation_steps=8,
+                gradient_checkpointing=True,
+                logging_dir=os.path.join(LOGS_DIR, f"run_{timestamp}"),
+                logging_steps=10,
+                save_strategy="epoch",
+                evaluation_strategy="epoch",
+                save_total_limit=2,
+                remove_unused_columns=False,  # Important for HF Spaces
+            )
+
+            dataset = Dataset.from_dict({
+                'text': [item['text'] for item in formatted_data]
+            })
+
+            data_collator = DataCollatorForLanguageModeling(
+                tokenizer=tokenizer,
+                mlm=False
+            )
+
+            # Custom progress callback
+            class ProgressCallback(gr.Progress):
+                def __init__(self, progress_callback, training_interface):
+                    self.progress_callback = progress_callback
+                    self.training_interface = training_interface
+                
+                def on_train_begin(self, args, state, control, **kwargs):
+                    if not self.training_interface.is_training:
+                        control.should_training_stop = True
+                    self.progress_callback(0.5, desc="Training started...")
+                
+                def on_epoch_begin(self, args, state, control, **kwargs):
+                    if not self.training_interface.is_training:
+                        control.should_training_stop = True
+                    epoch_progress = (state.epoch / args.num_train_epochs)
+                    total_progress = 0.5 + (epoch_progress * 0.4)
+                    self.progress_callback(total_progress, 
+                        desc=f"Training epoch {state.epoch + 1}/{args.num_train_epochs}...")
+
+            trainer = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=dataset,
+                data_collator=data_collator,
+                callbacks=[ProgressCallback(progress, self)]
+            )
+            
+            if not self.is_training:
+                return "Training cancelled."
+                
+            trainer.train()
+            
+            if not self.is_training:
+                return "Training cancelled."
+                
+            # Save model
+            progress(0.9, desc="Saving model...")
+            trainer.save_model()
+            tokenizer.save_pretrained(specific_output_dir)
+            
+            progress(1.0, desc="Training completed!")
+            return f"Training completed! Model saved in {specific_output_dir}"
+            
+        except Exception as e:
+            self.is_training = False
+            raise gr.Error(f"Training error: {str(e)}")
+
+def create_training_interface():
+    """Create Gradio interface"""
+    interface = TrainingInterface()
+    
+    with gr.Blocks(title="DeepSeek Model Training Interface") as app:
+        gr.Markdown("# DeepSeek Model Fine-tuning Interface")
+        
+        with gr.Row():
+            with gr.Column():
+                learning_rate = gr.Slider(
+                    minimum=1e-5,
+                    maximum=1e-3,
+                    value=2e-4,
+                    label="Learning Rate"
+                )
+                num_epochs = gr.Slider(
+                    minimum=1,
+                    maximum=10,
+                    value=3,
+                    step=1,
+                    label="Number of Epochs"
+                )
+                batch_size = gr.Slider(
+                    minimum=1,
+                    maximum=8,
+                    value=4,
+                    step=1,
+                    label="Batch Size"
+                )
+                
+        with gr.Row():
+            train_button = gr.Button("Start Training", variant="primary")
+            stop_button = gr.Button("Stop Training", variant="secondary")
+            
+        output_text = gr.Textbox(
+            label="Training Status",
+            placeholder="Training status will appear here...",
+            lines=10
+        )
+        
+        train_button.click(
+            fn=interface.train_model,
+            inputs=[learning_rate, num_epochs, batch_size],
+            outputs=output_text
+        )
+        
+        stop_button.click(
+            fn=interface.stop_training,
+            inputs=[],
+            outputs=output_text
+        )
 
+    return app
 
 if __name__ == "__main__":
-    demo.launch()
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    os.makedirs(LOGS_DIR, exist_ok=True)
+    
+    app = create_training_interface()
+    app.launch()