Gguf Inference

Source vllm-project/vllm.

  1.  1from huggingface_hub import hf_hub_download
  2.  2
  3.  3from vllm import LLM, SamplingParams
  4.  4
  5.  5
  6.  6def run_gguf_inference(model_path):
  7.  7    PROMPT_TEMPLATE = "<|system|>\n{system_message}</s>\n<|user|>\n{prompt}</s>\n<|assistant|>\n"  # noqa: E501
  8.  8    system_message = "You are a friendly chatbot who always responds in the style of a pirate."  # noqa: E501
  9.  9    # Sample prompts.
  10. 10    prompts = [
  11. 11        "How many helicopters can a human eat in one sitting?",
  12. 12        "What's the future of AI?",
  13. 13    ]
  14. 14    prompts = [
  15. 15        PROMPT_TEMPLATE.format(system_message=system_message, prompt=prompt)
  16. 16        for prompt in prompts
  17. 17    ]
  18. 18    # Create a sampling params object.
  19. 19    sampling_params = SamplingParams(temperature=0, max_tokens=128)
  20. 20
  21. 21    # Create an LLM.
  22. 22    llm = LLM(model=model_path,
  23. 23              tokenizer="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
  24. 24              gpu_memory_utilization=0.95)
  25. 25
  26. 26    outputs = llm.generate(prompts, sampling_params)
  27. 27    # Print the outputs.
  28. 28    for output in outputs:
  29. 29        prompt = output.prompt
  30. 30        generated_text = output.outputs[0].text
  31. 31        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
  32. 32
  33. 33
  34. 34if __name__ == "__main__":
  35. 35    repo_id = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
  36. 36    filename = "tinyllama-1.1b-chat-v1.0.Q4_0.gguf"
  37. 37    model = hf_hub_download(repo_id, filename=filename)
  38. 38    run_gguf_inference(model)