import torch
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer

def run_inference(model_id, prompt, max_tokens=50):
    print(f"Fetching model '{model_id}' from Hugging Face Hub...")
    
    # Load tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id, 
        torch_dtype=torch.float16, 
        device_map="auto"
    )
    
    print(f"\nPrompt: {prompt}")
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Llama models do not use token_type_ids
    inputs.pop("token_type_ids", None)

    # Generate
    print("Generating...")
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=max_tokens, 
            do_sample=True, 
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id
        )
    
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"\nResponse:\n{result}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_id", type=str, default="uisikdag/umitllama04b-galore-english")
    parser.add_argument("--prompt", type=str, default="The future of artificial intelligence is")
    args = parser.parse_args()
    
    run_inference(args.model_id, args.prompt)