infer_ifttt.py

"""
    # Inferenza singola: uso normale

    # Inferenza batch da CSV: aggiungere --input_csv test_inputs.csv --input_col natural_language --output_csv results.csv

    # Inferenza interattiva: aggiungere --interactive
"""

import argparse
import csv
import os
import sys
from typing import List, Optional

import pandas as pd
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


class IFTTTInferenceEngine:
    def __init__(
            self,
            model_path: str,
            device: str = "auto",
            input_prefix: str = "nl2js: "
    ):
        """
        Inizializza il motore di inferenza.

        Args:
            model_path: Percorso al modello addestrato
            device: Dispositivo ('cpu', 'cuda', 'mps', 'auto')
            input_prefix: Prefisso per l'input (deve corrispondere al training)
        """
        self.model_path = model_path
        self.input_prefix = input_prefix

        # Auto-detect device se richiesto
        if device == "auto":
            if torch.cuda.is_available():
                device = "cuda"
            elif torch.backends.mps.is_available():
                device = "mps"
            else:
                device = "cpu"

        self.device = device
        print(f"🚀 Caricando modello da: {model_path}")
        print(f"📱 Dispositivo: {device}")

        # Carica tokenizer e modello
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
            self.model.to(device)
            self.model.eval()
            print("✅ Modello caricato con successo!")
        except Exception as e:
            print(f"❌ Errore nel caricamento del modello: {e}")
            sys.exit(1)

    def generate(
            self,
            inputs: List[str],
            max_length: int = 124,
            num_beams: int = 4,
            temperature: float = 1.0,
            do_sample: bool = False,
            top_p: float = 0.9,
            batch_size: int = 4
    ) -> List[str]:
        """
        Genera codice JS per una lista di input in linguaggio naturale.

        Args:
            inputs: Lista di frasi in linguaggio naturale
            max_length: Lunghezza massima output
            num_beams: Numero di beam per beam search
            temperature: Temperatura per sampling (se do_sample=True)
            do_sample: Se usare sampling invece di beam search
            top_p: Nucleus sampling threshold
            batch_size: Dimensione batch per elaborazione

        Returns:
            Lista di codici JavaScript generati
        """
        all_outputs = []

        # Processa a batch per gestire memoria
        for i in range(0, len(inputs), batch_size):
            batch_inputs = inputs[i:i + batch_size]

            # Aggiungi prefisso task
            formatted_inputs = [self.input_prefix + inp for inp in batch_inputs]

            # Tokenizza
            with torch.no_grad():
                encoded = self.tokenizer(
                    formatted_inputs,
                    padding=True,
                    truncation=True,
                    max_length=124,
                    return_tensors="pt"
                ).to(self.device)

                # Genera
                if do_sample:
                    # Sampling mode
                    generated = self.model.generate(
                        **encoded,
                        max_length=max_length,
                        temperature=temperature,
                        do_sample=True,
                        top_p=top_p,
                        pad_token_id=self.tokenizer.pad_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                        early_stopping=True
                    )
                else:
                    # Beam search mode (default)
                    generated = self.model.generate(
                        **encoded,
                        max_length=max_length,
                        num_beams=num_beams,
                        pad_token_id=self.tokenizer.pad_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                        early_stopping=True
                    )

                # Decodifica
                batch_outputs = self.tokenizer.batch_decode(
                    generated,
                    skip_special_tokens=True
                )

                all_outputs.extend([out.strip() for out in batch_outputs])

                # Progress
                # print(f"Elaborati {min(i + batch_size, len(inputs))}/{len(inputs)} input")

        return all_outputs

    def single_inference(self, input_text: str, **kwargs) -> str:
        """Inferenza singola - wrapper conveniente."""
        return self.generate([input_text], **kwargs)[0]


def load_input_csv(csv_path: str, input_col: str) -> List[str]:
    """Carica input da CSV."""
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"File non trovato: {csv_path}")

    df = pd.read_csv(csv_path)
    if input_col not in df.columns:
        # Prova a indovinare la colonna
        text_cols = [c for c in df.columns if df[c].dtype == object]
        if text_cols:
            input_col = text_cols[0]
            print(f"⚠️  Colonna '{input_col}' non trovata, uso '{input_col}'")
        else:
            raise ValueError(f"Colonna '{input_col}' non trovata nel CSV")

    inputs = df[input_col].dropna().astype(str).tolist()
    print(f"📄 Caricati {len(inputs)} input da {csv_path}")
    return inputs


def save_results_csv(
        inputs: List[str],
        outputs: List[str],
        output_path: str,
        input_col_name: str = "input",
        output_col_name: str = "generated_js"
):
    """Salva risultati in CSV."""
    df = pd.DataFrame({
        input_col_name: inputs,
        output_col_name: outputs
    })
    df.to_csv(output_path, index=False, encoding='utf-8')
    print(f"💾 Risultati salvati in: {output_path}")


def interactive_mode(engine: IFTTTInferenceEngine):
    """Modalità interattiva."""
    print("\n🎯 Modalità interattiva IFTTT Filter Code Generator")
    print("Digita frasi in linguaggio naturale per generare codice JavaScript.")
    print("Comandi speciali:")
    print("  - 'quit' o 'exit' per uscire")
    print("  - 'help' per mostrare esempi")
    print("  - 'config' per cambiare parametri generazione")
    print("-" * 60)

    # Parametri di default
    params = {
        'num_beams': 4,
        'do_sample': False,
        'temperature': 1.0,
        'top_p': 0.9
    }

    while True:
        try:
            user_input = input("\n🔤 Input: ").strip()

            if user_input.lower() in ['quit', 'exit', 'q']:
                print("👋 Arrivederci!")
                break

            elif user_input.lower() == 'help':
                print("\n📝 Esempi di input:")
                print("  - 'non deve quando piove'")
                print("  - 'solo di mattina'")
                print("  - 'se l'email contiene promo'")
                print("  - 'mai quando c'è il sole'")
                print("  - 'esclusivamente nei weekend'")
                continue

            elif user_input.lower() == 'config':
                print(f"\n⚙️  Parametri attuali:")
                for k, v in params.items():
                    print(f"  {k}: {v}")
                print("\nPremi ENTER per mantenere il valore attuale")

                try:
                    params['num_beams'] = int(input("num_beams (2-8): ") or params['num_beams'])
                    sample = input("do_sample (true/false): ").lower()
                    if sample in ['true', '1', 'yes']:
                        params['do_sample'] = True
                        params['temperature'] = float(input("temperature (0.1-2.0): ") or params['temperature'])
                        params['top_p'] = float(input("top_p (0.1-1.0): ") or params['top_p'])
                    else:
                        params['do_sample'] = False
                except ValueError:
                    print("⚠️  Valori non validi, uso default")
                continue

            elif not user_input:
                continue

            # Genera
            # print("🔄 Generando...")
            result = engine.single_inference(user_input, **params)
            print(f"🔧 Output: {result}")

        except KeyboardInterrupt:
            print("\n👋 Interrotto dall'utente. Arrivederci!")
            break
        except Exception as e:
            print(f"❌ Errore: {e}")


def main():
    parser = argparse.ArgumentParser(description="Inferenza modello IFTTT NL->JS")

    # Modello
    parser.add_argument("--model_path", default='flan_base_ifttt_out', help="Percorso al modello addestrato")
    parser.add_argument("--device", default="auto", choices=["cpu", "cuda", "mps", "auto"])
    parser.add_argument("--input_prefix", default="nl2js: ", help="Prefisso input (deve corrispondere al training)")

    # Input modes
    parser.add_argument("--input", type=str, help="Singolo input da processare")
    parser.add_argument("--input_csv", type=str, help="CSV con input da processare")
    parser.add_argument("--input_col", default="natural_language", help="Colonna input nel CSV")
    parser.add_argument("--interactive", action="store_true", help="Modalità interattiva")

    # Output
    parser.add_argument("--output_csv", type=str, help="File CSV output (obbligatorio se --input_csv)")

    # Parametri generazione
    parser.add_argument("--max_length", type=int, default=124)
    parser.add_argument("--num_beams", type=int, default=4)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--do_sample", action="store_true")
    parser.add_argument("--top_p", type=float, default=0.9)

    args = parser.parse_args()

    # Validazione argomenti
    if not any([args.input, args.input_csv, args.interactive]):
        parser.error("Specificare almeno uno tra: --input, --input_csv, --interactive")

    if args.input_csv and not args.output_csv:
        parser.error("--output_csv è obbligatorio quando si usa --input_csv")

    # Inizializza engine
    engine = IFTTTInferenceEngine(
        model_path=args.model_path,
        device=args.device,
        input_prefix=args.input_prefix
    )

    # Parametri generazione
    gen_params = {
        'max_length': args.max_length,
        'num_beams': args.num_beams,
        'temperature': args.temperature,
        'do_sample': args.do_sample,
        'top_p': args.top_p,
        'batch_size': args.batch_size
    }

    try:
        if args.interactive:
            # Modalità interattiva
            interactive_mode(engine)

        elif args.input:
            # Singolo input
            print(f"🔤 Input: {args.input}")
            result = engine.single_inference(args.input, **gen_params)
            print(f"🔧 Output: {result}")

        elif args.input_csv:
            # Batch da CSV
            inputs = load_input_csv(args.input_csv, args.input_col)
            print(f"🔄 Processando {len(inputs)} input...")

            outputs = engine.generate(inputs, **gen_params)

            save_results_csv(
                inputs=inputs,
                outputs=outputs,
                output_path=args.output_csv,
                input_col_name=args.input_col,
                output_col_name="generated_js"
            )

            # Mostra alcuni esempi
            print(f"\n📋 Primi 3 risultati:")
            for i, (inp, out) in enumerate(zip(inputs[:3], outputs[:3])):
                print(f"  {i + 1}. '{inp}' → '{out}'")

    except Exception as e:
        print(f"❌ Errore durante l'inferenza: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()