Koios-API/KoiosAPI-codegemma-7b-it

1import torch2from transformers import AutoModelForCausalLM, AutoTokenizer, GemmaTokenizer, AutoConfig, TrainingArguments, Trainer, BertLMHeadModel, BertForSequenceClassification3from datasets import Dataset4import pandas as pd5import csv6from transformers import TrainingArguments, Trainer7 8import tensorflow as tf9# Check TensorFlow GPU availability10print("GPUs Available: ", tf.config.list_physical_devices('GPU'))11 12import os13# Setting the environment variable for MPS14os.environ['PYTORCH_MPS_HIGH_WATERMARK_RATIO'] = '0.0'15 16def get_device():17    """Automatically chooses the best device."""18    if torch.cuda.is_available():19        return torch.device('cuda')20    elif torch.backends.mps.is_available():21        return torch.device('mps')22    else:23        return torch.device('cpu')24 25def load_data_and_config(data_path):26    """Loads training data from CSV."""27    data = []28    with open(data_path, newline='', encoding='utf-8') as csvfile:29        reader = csv.DictReader(csvfile, delimiter=';')30        for row in reader:31            data.append({'text': row['description']})32    return data33 34def train_model(model, tokenizer, data, device):35    """Trains the model using the Hugging Face Trainer API."""36    # inputs = [tokenizer(d['text'], max_length=256, truncation=True, padding='max_length', return_tensors="pt") for d in data]37    inputs = [tokenizer(d['text'], max_length=256, truncation=True, padding='max_length', return_tensors="pt").to(torch.float16) for d in data]38    dataset = Dataset.from_dict({39        'input_ids': [x['input_ids'].squeeze() for x in inputs],40        'labels': [x['input_ids'].squeeze() for x in inputs]41    })42    43    training_args = TrainingArguments(44        output_dir='./results',45        num_train_epochs=3,46        per_device_train_batch_size=1,47        gradient_accumulation_steps=4,48        fp16=True,  # Enable mixed precision49        warmup_steps=500,50        weight_decay=0.01,51        logging_dir='./logs',52        logging_steps=10,53    )54    55    trainer = Trainer(56        model=model,57        args=training_args,58        train_dataset=dataset,59        tokenizer=tokenizer60    )61    62    trainer.train()63 64     # Optionally clear cache if using GPU or MPS65    if torch.cuda.is_available():66        print(torch.cuda.memory_summary(device=None, abbreviated=False))67        torch.cuda.empty_cache()68    elif torch.has_mps:69        torch.mps.empty_cache()70 71    # Perform any remaining steps such as logging, saving, etc.72    trainer.save_model()73 74def main(api_name, base_url):75    device = get_device()  # Get the appropriate device76    data = load_data_and_config("train2.csv")77 78    model_id = "google/codegemma-2b"79    tokenizer = GemmaTokenizer.from_pretrained(model_id)80    model = AutoModelForCausalLM.from_pretrained(model_id)81 82    #tokenizer = AutoTokenizer.from_pretrained("google/codegemma-2b")83    # Load the configuration for a specific model84    # config = AutoConfig.from_pretrained('google/codegemma-2b')85    # Update the activation function86    # config.hidden_act = ''  # Set to use approximate GeLU gelu_pytorch_tanh87    # config.hidden_activation = 'gelu_pytorch_tanh'  # Set to use GeLU88 89    # model = AutoModelForCausalLM.from_pretrained('google/codegemma-2b', is_decoder=True)90    #model = BertLMHeadModel.from_pretrained('google/codegemma-2b', is_decoder=True)91    # Example assuming you have a prepared dataset for classification92    #model = BertForSequenceClassification.from_pretrained('thenlper/gte-small', num_labels=2, is_decoder=True)  # binary classification93    94    model.to(device)  # Move model to the appropriate device95 96    train_model(model, tokenizer, data, device)97    98    model.save_pretrained("./fine_tuned_model")99    tokenizer.save_pretrained("./fine_tuned_model")100 101    prompt = "I need to retrieve the latest block on chain using a python script"102    api_query = generate_api_query(model, tokenizer, prompt, "latest block on chain", api_name, base_url)103    print(f"Generated code: {api_query}")104 105def generate_api_query(model, tokenizer, prompt, desired_output, api_name, base_url):106    # Prepare input prompt for the model, ensure tensors are compatible with PyTorch107    input_ids = tokenizer.encode(f"{prompt} Write an API query to {api_name} to get {desired_output}", return_tensors="pt")108 109    # Ensure input_ids are on the same device as the model110    input_ids = input_ids.to(model.device)111 112    # Generate query using model with temperature for randomness113    output = model.generate(input_ids, max_length=128, temperature=0.001, do_sample=True)114    115 116    # Decode the generated query tokens117    query = tokenizer.decode(output[0], skip_special_tokens=True)118    return f"{base_url}/{query}"119 120if __name__ == "__main__":121    api_name = "Koios"122    base_url = "https://api.koios.rest/v1"123    main(api_name, base_url)124