commit from tesla

app.py

@@ -33,7 +33,8 @@ def load_data_and_config(data_path):
 
 def train_model(model, tokenizer, data, device):
     """Trains the model using the Hugging Face Trainer API."""
-    inputs = [tokenizer(d['text'], max_length=256, truncation=True, padding='max_length', return_tensors="pt") for d in data]
+    # inputs = [tokenizer(d['text'], max_length=256, truncation=True, padding='max_length', return_tensors="pt") for d in data]
+    inputs = [tokenizer(d['text'], max_length=256, truncation=True, padding='max_length', return_tensors="pt").to(torch.float16) for d in data]
     dataset = Dataset.from_dict({
         'input_ids': [x['input_ids'].squeeze() for x in inputs],
         'labels': [x['input_ids'].squeeze() for x in inputs]
@@ -43,7 +44,7 @@ def train_model(model, tokenizer, data, device):
         output_dir='./results',
         num_train_epochs=3,
         per_device_train_batch_size=1,
-        gradient_accumulation_steps=2,
+        gradient_accumulation_steps=4,
         fp16=True,  # Enable mixed precision
         warmup_steps=500,
         weight_decay=0.01,
@@ -61,9 +62,8 @@ def train_model(model, tokenizer, data, device):
     trainer.train()
 
      # Optionally clear cache if using GPU or MPS
-    print(torch.cuda.memory_summary(device=None, abbreviated=False))
-
     if torch.cuda.is_available():
+        print(torch.cuda.memory_summary(device=None, abbreviated=False))
         torch.cuda.empty_cache()
     elif torch.has_mps:
         torch.mps.empty_cache()
@@ -85,6 +85,9 @@ def main(api_name, base_url):
     #model = BertLMHeadModel.from_pretrained('google/codegemma-2b', is_decoder=True)
     # Example assuming you have a prepared dataset for classification
     #model = BertForSequenceClassification.from_pretrained('thenlper/gte-small', num_labels=2, is_decoder=True)  # binary classification
+    # Example: Offloading embeddings to CPU
+    model.embeddings.to('cpu')
+
     model.to(device)  # Move model to the appropriate device
 
     train_model(model, tokenizer, data, device)
@@ -104,7 +107,8 @@ def generate_api_query(model, tokenizer, prompt, desired_output, api_name, base_
     input_ids = input_ids.to(model.device)
 
     # Generate query using model with temperature for randomness
-    output = model.generate(input_ids, max_length=128, truncation=True, padding='max_length', temperature=0.1, do_sample=True)
+    output = model.generate(input_ids, max_length=128, temperature=0.001, do_sample=True)
+    
 
     # Decode the generated query tokens
     query = tokenizer.decode(output[0], skip_special_tokens=True)