commit from tesla

app.py

@@ -42,8 +42,9 @@ def train_model(model, tokenizer, data, device):
     training_args = TrainingArguments(
         output_dir='./results',
         num_train_epochs=3,
-        per_device_train_batch_size=1,
-        gradient_accumulation_steps=1,
+        per_device_train_batch_size=8,
+        gradient_accumulation_steps=4,
+        fp16=True,  # Enable mixed precision
         warmup_steps=500,
         weight_decay=0.01,
         logging_dir='./logs',
@@ -89,7 +90,7 @@ 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=256, temperature=0.1, do_sample=True)
+    output = model.generate(input_ids, max_length=128, truncation=True, padding='max_length', temperature=0.1, do_sample=True)
 
     # Decode the generated query tokens
     query = tokenizer.decode(output[0], skip_special_tokens=True)