Skip to content

Instantly share code, notes, and snippets.

@john-adeojo

john-adeojo/haystack_vllm.py

Created January 30, 2024 13:06
Show Gist options
  • Save john-adeojo/85200776d36184e446dd9ad33c49f955 to your computer and use it in GitHub Desktop.
Save john-adeojo/85200776d36184e446dd9ad33c49f955 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
haystack_vllm.py
from vllm_haystack import vLLMInvocationLayer
from haystack.nodes import PromptNode, PromptModel
from haystack.agents.conversational import ConversationalAgent
llm = PromptModel(model_path_or_path=model_path, invocation_layer_class=vLLMInvocationLayer, api_key="EMPTY", model_kwargs={
"api_base" : vllm_endpoint,
"maximum_context_length": 250,
})
prompt_node = PromptNode(llm)
prompt_node("What is AI?")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment