2

app.py

@@ -4,6 +4,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments,
 import csv
 import yaml
 from datasets import Dataset
+
+import tensorflow as tf
+# Check TensorFlow GPU availability
+print("GPUs Available: ", tf.config.list_physical_devices('GPU'))
+
 import os
 os.environ['PYTORCH_MPS_HIGH_WATERMARK_RATIO'] = '0.0'
 
@@ -19,7 +24,8 @@ def load_data_and_config(data_path):
 def generate_api_query(model, tokenizer, prompt, desired_output, api_name, base_url):
     """Generates an API query using a fine-tuned model."""
     input_ids = tokenizer.encode(prompt + f" Write an API query to {api_name} to get {desired_output}", return_tensors="pt")
-    output = model.generate(input_ids, max_length=256, temperature=0.7)
+    input_ids = input_ids.to(model.device)  # Ensure input_ids are on the same device as the model
+    output = model.generate(input_ids, max_length=256, temperature=0.7, do_sample=True)  # Enable sampling with temperature control
     query = tokenizer.decode(output[0], skip_special_tokens=True)
     return f"{base_url}/{query}"
 
@@ -58,19 +64,19 @@ def train_model(model, tokenizer, data):
     # Optionally clear cache if using GPU or MPS
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    elif torch.has_mps:
+    elif torch.backends.mps.is_built():
         torch.mps.empty_cache()
 
     # Perform any remaining steps such as logging, saving, etc.
     trainer.save_model()
 
-if __name__ == "__main__":
-    # Load data and configurations
+def main(api_name, base_url):
+    # Load data
     data = load_data_and_config("train2.csv")
 
     # Load tokenizer and model
-    tokenizer = AutoTokenizer.from_pretrained("google/codegemma-7b-it")
-    model = AutoModelForCausalLM.from_pretrained("google/codegemma-7b-it")
+    tokenizer = AutoTokenizer.from_pretrained("thenlper/gte-small")
+    model = AutoModelForCausalLM.from_pretrained("thenlper/gte-small")
 
     # Train the model on your dataset
     train_model(model, tokenizer, data)
@@ -81,5 +87,10 @@ if __name__ == "__main__":
 
     # Example usage
     prompt = "I need to retrieve the latest block on chain using a python script"
-    api_query = generate_api_query(model, tokenizer, prompt, "latest block on chain", config["api_name"], config["base_url"])
+    api_query = generate_api_query(model, tokenizer, prompt, "latest block on chain", api_name, base_url)
     print(f"Generated code: {api_query}")
+
+if __name__ == "__main__":
+    api_name = "Koios"
+    base_url = "https://api.koios.rest"
+    main(api_name, base_url)