Add and verify LLM training and inference examples

remote-gpu-client/examples/inference_llm.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""
+Remote LLM Inference
+====================
+Usa il modello LLM salvato sul server (nella directory persistente)
+per generare testo da un prompt.
+"""
+
+import sys
+from antigravity_sdk import RemoteGPU
+
+# Default prompts
+PROMPTS = [
+    "The future of artificial intelligence is",
+    "Once upon a time in a digital world,",
+    "The best way to learn Python is"
+]
+
+INFERENCE_CODE = r'''
+import os
+import sys
+
+print("🔧 Setting up Environment...")
+os.system(f"{sys.executable} -m pip install transformers==4.37.2 accelerate==0.27.2 --quiet")
+
+
+
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# 1. Configuration
+STORAGE_DIR = "/home/user/app/storage/my_llm"
+print(f"📂 Loading LLM from {STORAGE_DIR}...")
+
+if not os.path.exists(STORAGE_DIR):
+    print("❌ Model not found! Run 'train_llm.py' first.")
+    sys.exit(1)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"   Using device: {device}")
+
+# 2. Load Model & Tokenizer
+try:
+    tokenizer = AutoTokenizer.from_pretrained(STORAGE_DIR)
+    model = AutoModelForCausalLM.from_pretrained(STORAGE_DIR).to(device)
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    sys.exit(1)
+
+# 3. Generate
+prompts = {PROMPTS_PLACEHOLDER}
+
+print("\n🔮 Generating Text...")
+print("=" * 60)
+
+for prompt in prompts:
+    print(f"📝 Prompt: {prompt}")
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        output = model.generate(
+            **inputs, 
+            max_length=100, 
+            num_return_sequences=1,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    print(f"🤖 Output:\n{generated_text}")
+    print("-" * 60)
+
+print("✅ Generation Complete.")
+'''
+
+def main():
+    print("📡 Connecting to Remote GPU for LLM Inference...")
+    gpu = RemoteGPU()
+    
+    # Inject prompts into code
+    code_to_run = INFERENCE_CODE.replace("{PROMPTS_PLACEHOLDER}", str(PROMPTS))
+    
+    result = gpu.run(code_to_run)
+
+if __name__ == "__main__":
+    main()

remote-gpu-client/examples/train_llm.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""
+Remote LLM Training (DistilGPT2 on Wikitext)
+============================================
+Allena un piccolo LLM (DistilGPT2) su un dataset di testo (Wikitext-2)
+direttamente sulla GPU remota e salva il modello persistente.
+"""
+
+from antigravity_sdk import RemoteGPU
+
+TRAINING_CODE = r'''
+import os
+import sys
+
+print("🔧 Setting up Environment...")
+# Pin compatible versions for PyTorch 2.1.2
+os.system(f"{sys.executable} -m pip install transformers==4.37.2 datasets==2.17.0 accelerate==0.27.2 --quiet")
+
+
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, TextDataset, DataCollatorForLanguageModeling
+from datasets import load_dataset
+
+print("🚀 Starting LLM Training...")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"   Using device: {device}")
+
+# 1. Configuration
+MODEL_NAME = "distilgpt2"
+STORAGE_DIR = "/home/user/app/storage/my_llm"
+os.makedirs(STORAGE_DIR, exist_ok=True)
+
+# 2. Load Tokenizer & Model
+print(f"   Loading {MODEL_NAME}...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer.pad_token = tokenizer.eos_token # Fix for GPT-2 which has no pad token
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
+
+
+# 3. Prepare Dataset (Wikitext-2 small subset for speed)
+print("   Loading dataset (wikitext-2)...")
+# For simplicity/speed in this demo, accessing a small raw text subset or using 'wikitext' library
+dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1%]") # 1% just for demo speed
+print(f"   Dataset loaded. Rows: {len(dataset)}")
+
+# Helper to tokenize
+def tokenize_function(examples):
+    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
+
+print("   Tokenizing...")
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+tokenized_datasets = tokenized_datasets.remove_columns(["text"])
+tokenized_datasets.set_format("torch")
+
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+# 4. Training Arguments
+training_args = TrainingArguments(
+    output_dir="./results",
+    overwrite_output_dir=True,
+    num_train_epochs=1,
+    per_device_train_batch_size=4,
+    save_steps=500,
+    save_total_limit=1,
+    report_to="none",
+    disable_tqdm=True # Cleaner output logs
+)
+
+# 5. Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets,
+    data_collator=data_collator,
+)
+
+# 6. Train
+print("   Starting Fine-Tuning...")
+trainer.train()
+
+# 7. Save Persistently
+print(f"   💾 Saving model to {STORAGE_DIR}...")
+model.save_pretrained(STORAGE_DIR)
+tokenizer.save_pretrained(STORAGE_DIR)
+
+# 8. Test Generation
+print("   Testing generation...")
+input_text = "The future of AI is"
+inputs = tokenizer(input_text, return_tensors="pt").to(device)
+output = model.generate(**inputs, max_length=50, num_return_sequences=1)
+generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+
+print("-" * 40)
+print(f"Input: {input_text}")
+print(f"Output: {generated_text}")
+print("-" * 40)
+
+print("✅ LLM Training Complete & Model Saved.")
+'''
+
+def main():
+    print("📡 Connecting to Remote GPU for LLM Training...")
+    gpu = RemoteGPU()
+    
+    # Run, getting logs mostly
+    result = gpu.run(TRAINING_CODE)
+    
+    if "Training Complete" in result.output:
+        print("\n🏆 LLM Addestrato e Salvato sul Server!")
+    else:
+        print("\n⚠️ Qualcosa è andato storto (controlla i log sopra).")
+
+if __name__ == "__main__":
+    main()