from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Download the GGUF model
model_path = hf_hub_download(
    repo_id="Habibur2/Phi-3.5-mini-GGUF",
    filename="phi-3.5-mini-q4_0.gguf"
)

# Load the model
# Set n_gpu_layers=-1 for full GPU usage (Requires CUDA)
# Set n_gpu_layers=0 if you only want to use CPU
llm = Llama(
    model_path=model_path,
    n_ctx=2048,        # Context window
    n_threads=4,       # Number of CPU threads
    n_gpu_layers=-1    # Offload all layers to GPU
)

# Run Inference
output = llm.create_chat_completion(
    messages=[
        {"role": "system", "content": "You are a helpful AI assistant."},
        {"role": "user", "content": "Who is the founder of Microsoft?"}
    ],
    max_tokens=512,
    temperature=0.7
)

print(output['choices'][0]['message']['content'])