query_llm.py

import argparse
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import os
import time

# Config file name
CONFIG_FILE = "model_path.cfg"

def get_default_model_path():
    if os.path.exists(CONFIG_FILE):
        with open(CONFIG_FILE, "r") as f:
            return f.read().strip()
    return None

def main():
    default_path = get_default_model_path()
    
    parser = argparse.ArgumentParser(description="Query a local Llama 7B model.")
    parser.add_argument("prompt", type=str, help="The text prompt to send to the LLM.")
    parser.add_argument("--model-path", type=str, default=default_path, help="Path to the local model directory. Defaults to path in model_path.cfg")
    parser.add_argument("--cpu", action="store_true", help="Force the model to run on CPU instead of GPU.")
    parser.add_argument("--max-new-tokens", type=int, default=100, help="Maximum number of new tokens to generate.")
    
    args = parser.parse_args()

    if not args.model_path:
        print(f"Error: No model path specified and {CONFIG_FILE} not found or empty.")
        print(f"Please create {CONFIG_FILE} with the model path or use --model-path argument.")
        return

    print(f"Loading model from: {args.model_path}")
    
    # Determine device and loading arguments
    if args.cpu:
        print("Running on CPU.")
        device_map = None
        dtype = torch.float32
    else:
        if torch.cuda.is_available():
            print(f"Running on GPU: {torch.cuda.get_device_name(0)}")
            device_map = "auto"
            dtype = torch.float16
        else:
            print("CUDA not available. Falling back to CPU.")
            device_map = None
            dtype = torch.float32

    try:
        # Load Tokenizer
        print("Loading tokenizer...")
        tokenizer = AutoTokenizer.from_pretrained(args.model_path, local_files_only=True)
        
        # Load Model
        print("Loading model (this may take a while)...")
        model = AutoModelForCausalLM.from_pretrained(
            args.model_path,
            torch_dtype=dtype,
            local_files_only=True,
            device_map=device_map,
        )
        
        # If we are on CPU (device_map is None), we don't need to move it manually 
        # because it loads to CPU by default. 
        # If we are on GPU (device_map="auto"), accelerate handles it.


        print("Model loaded successfully.")

        # Prepare input
        inputs = tokenizer(args.prompt, return_tensors="pt").to(model.device)

        # Generate
        print("Generating response...")
        start_time = time.time()
        with torch.no_grad():
            outputs = model.generate(
                **inputs, 
                max_new_tokens=args.max_new_tokens,
                do_sample=True,
                temperature=0.7,
                top_p=0.9
            )
        end_time = time.time()

        # Decode output
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        print("-" * 50)
        print(f"Response ({end_time - start_time:.2f}s):")
        print(response)
        print("-" * 50)

    except Exception as e:
        print(f"An error occurred: {e}")
        print("\nPlease ensure the model path is correct and you have the necessary libraries installed.")

if __name__ == "__main__":
    main()