Upload handler.py

handler.py

@@ -0,0 +1,200 @@
+#python getvars-generic.py 1 --train_data seq_train_truck_distance.txt --model_output ./json_extraction_truck_distance
+#python getvars-generic.py 1 --train_data seq_train_data_add_point.txt --model_output ./json_extraction_add_point
+#python getvars-generic.py 1 --train_data seq_train_data_add_point_with_alerts.txt --model_output ./json_extraction_add_point_with_alerts
+
+#python getvars-generic.py 1 --train_data seq_train_point_activity.txt --model_output ./json_extraction_point_activity
+
+
+import torch
+import argparse
+import json
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+from datasets import Dataset
+from transformers import TrainingArguments, Trainer
+import random
+import json
+from huggingface_hub import HfApi, login, upload_folder, create_repo
+import os
+import gradio as gr
+
+from transformers import BitsAndBytesConfig
+
+# read variables from config.json file
+
+with open('config.json', 'r') as config_file:
+    config = json.load(config_file) 
+
+num_args = len(config)
+
+
+mode = config.get('mode', '1') 
+train_data_path = config.get('train_data', 'seq_train_point_activity.txt') 
+model_output_path = config.get('model_output', 'json_extraction_point_activity')
+
+print(f"current mode: {mode}");
+print(f"trin data path: {train_data_path}");
+print(f"model output path: {model_output_path}");
+
+def train_model(train_data_path, model_output_path):
+    # Read JSON data from the provided file
+    with open(train_data_path, "r", encoding="utf-8") as file:
+        train_data = json.load(file)
+    
+    # Shuffle data randomly
+    random.shuffle(train_data)
+
+    # Define split ratio (95% train, 5% eval)
+    split_ratio = 0.95
+    split_index = int(len(train_data) * split_ratio)
+
+    # Split into training and evaluation sets
+    train_set = train_data[:split_index]
+    eval_set = train_data[split_index:]
+
+    # Create Hugging Face Dataset objects
+    train_dataset = Dataset.from_dict({
+        "input_text": [x["input"] for x in train_set],
+        "target_text": [x["output"] for x in train_set]
+    })
+
+    eval_dataset = Dataset.from_dict({
+        "input_text": [x["input"] for x in eval_set],
+        "target_text": [x["output"] for x in eval_set]
+    })
+
+    tokenizer = T5Tokenizer.from_pretrained("t5-large")
+    model = T5ForConditionalGeneration.from_pretrained("t5-large")
+    
+    def tokenize_function(examples):
+        model_inputs = tokenizer(examples["input_text"], padding="max_length", truncation=True, max_length=128)
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(examples["target_text"], padding="max_length", truncation=True, max_length=128)
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+    
+    train_dataset = train_dataset.map(tokenize_function, batched=True)
+    eval_dataset = eval_dataset.map(tokenize_function, batched=True)
+    
+    training_args = TrainingArguments(
+        output_dir=model_output_path,
+        logging_dir="./logs",
+        logging_steps=10,
+        evaluation_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        num_train_epochs=35,
+        weight_decay=0.01
+    )
+    
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset
+    )
+    
+    print("Starting training...")
+    trainer.train()
+    
+    model.save_pretrained(model_output_path)
+    tokenizer.save_pretrained(model_output_path)
+
+    #for push repository
+    repo_name = "Reyad-Ahmmed/hf-data-timeframe" 
+
+    # Your repository name 
+    api_token = os.getenv("hf_token")  # Retrieve the API token from environment variable
+
+    if not api_token:
+        raise ValueError("API token not found. Please set the HF_API_TOKEN environment variable.")
+
+    # Create repository (if not already created)
+    api = HfApi()
+    create_repo(repo_id=repo_name, token=api_token, exist_ok=True)
+
+    # Upload the model and tokenizer to the Hugging Face repository
+    
+    upload_folder(
+        folder_path=f"{model_output_path}",
+        path_in_repo=f"{model_output_path}",
+        repo_id=repo_name,
+        token=api_token,
+        commit_message="Push getvar generic t5 model",
+        #overwrite=True  # Force overwrite existing files
+    )
+
+    upload_folder(
+        folder_path=f"{model_output_path}",
+        path_in_repo=f"{model_output_path}",
+        repo_id=repo_name,
+        token=api_token,
+        commit_message="Push getvar generic t5 tokenizer",
+        #overwrite=True  # Force overwrite existing files
+    )
+
+    
+    print(f"Model training complete and saved to {model_output_path}.")
+
+def inference_loop(model_output_path):
+
+    
+    # Load model and tokenizer
+    #model = T5ForConditionalGeneration.from_pretrained(model_output_path)
+    #tokenizer = T5Tokenizer.from_pretrained(model_output_path)
+
+    model_name = "Reyad-Ahmmed/hf-data-timeframe"
+
+    model = T5ForConditionalGeneration.from_pretrained(model_name, subfolder=model_output_path)
+    tokenizer = T5Tokenizer.from_pretrained(model_name, subfolder=model_output_path)
+
+    # Enable 8-bit or 4-bit quantization
+    quantization_config = BitsAndBytesConfig(   # Set `load_in_4bit=True` for 4-bit quantization
+        llm_int8_threshold=8.0,
+        load_in_4bit = True
+    )
+
+    # Load quantized model
+    model = T5ForConditionalGeneration.from_pretrained(
+        model_name, subfolder=model_output_path,
+        quantization_config=quantization_config,
+        device_map="auto"  # Automatically uses GPU if available
+    )
+    
+    # Function to generate JSON output
+    def generate_json(input_text):
+        input_text = "Extract structured JSON for: " + input_text
+        #input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+
+        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)
+        
+        output_ids = model.generate(input_ids, max_length=100, temperature=0.3)
+        json_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        
+        try:
+            return json.dumps(json.loads(json_output), indent=2)  # Pretty-print JSON
+        except:
+            return json_output  # Return as string if not valid JSON
+
+    # Launch Gradio Interface
+    iface = gr.Interface(fn=generate_json, inputs="text", outputs="text", title="JSON Extractor")
+    iface.launch(share=True)
+
+
+if __name__ == "__main__":
+    #parser = argparse.ArgumentParser(description="T5 JSON Extraction Script")
+    #parser.add_argument("mode", type=int, help="1 for training, 2 for inference")
+    #parser.add_argument("--train_data", type=str, help="Path to training data file", required=False)
+    #parser.add_argument("--model_output", type=str, help="Path to save/load model", required=True)
+    #args = parser.parse_args()
+    
+    if mode == 1:
+        if not train_data_path:
+            print("Training mode requires --train_data argument.")
+        else:
+            train_model(train_data_path, model_output_path)
+            inference_loop(model_output_path)
+    elif mode == 2:
+        inference_loop(model_output_path)
+    else:
+        print("Invalid mode. Use 1 for training and 2 for inference.")