| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline | |
| import os | |
| # Same output dir as train.py (works from any cwd) | |
| SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| model_path = os.path.join(SCRIPT_DIR, "multilingual-doc-model") | |
| if not os.path.isdir(model_path): | |
| print(f"Model not found at {model_path}. Run train.py first to train the model.") | |
| exit(1) | |
| tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| model = AutoModelForCausalLM.from_pretrained(model_path) | |
| # Use GPU if available, else CPU | |
| device = 0 if __import__("torch").cuda.is_available() else -1 | |
| pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device) | |
| prompt = """User: Explícame este documento: | |
| La IA mejora la productividad. | |
| Assistant:""" | |
| result = pipe(prompt, max_new_tokens=120, do_sample=True, temperature=0.7) | |
| print(result[0]["generated_text"]) | |