Harrison/self hosted runhouse (#1154)

Co-authored-by: Donny Greenberg <dongreenberg2@gmail.com>
Co-authored-by: John Dagdelen <jdagdelen@users.noreply.github.com>
Co-authored-by: Harrison Chase <harrisonchase@Harrisons-MBP.attlocal.net>
Co-authored-by: Andrew White <white.d.andrew@gmail.com>
Co-authored-by: Peng Qu <82029664+pengqu123@users.noreply.github.com>
Co-authored-by: Matt Robinson <mthw.wm.robinson@gmail.com>
Co-authored-by: jeff <tangj1122@gmail.com>
Co-authored-by: Harrison Chase <harrisonchase@Harrisons-MacBook-Pro.local>
Co-authored-by: zanderchase <zander@unfold.ag>
Co-authored-by: Charles Frye <cfrye59@gmail.com>
Co-authored-by: zanderchase <zanderchase@gmail.com>
Co-authored-by: Shahriar Tajbakhsh <sh.tajbakhsh@gmail.com>
Co-authored-by: Stefan Keselj <skeselj@princeton.edu>
Co-authored-by: Francisco Ingham <fpingham@gmail.com>
Co-authored-by: Dhruv Anand <105786647+dhruv-anand-aintech@users.noreply.github.com>
Co-authored-by: cragwolfe <cragcw@gmail.com>
Co-authored-by: Anton Troynikov <atroyn@users.noreply.github.com>
Co-authored-by: William FH <13333726+hinthornw@users.noreply.github.com>
Co-authored-by: Oliver Klingefjord <oliver@klingefjord.com>
Co-authored-by: blob42 <contact@blob42.xyz>
Co-authored-by: blob42 <spike@w530>
Co-authored-by: Enrico Shippole <henryshippole@gmail.com>
Co-authored-by: Ibis Prevedello <ibiscp@gmail.com>
Co-authored-by: jped <jonathanped@gmail.com>
Co-authored-by: Justin Torre <justintorre75@gmail.com>
Co-authored-by: Ivan Vendrov <ivan@anthropic.com>
Co-authored-by: Sasmitha Manathunga <70096033+mmz-001@users.noreply.github.com>
Co-authored-by: Ankush Gola <9536492+agola11@users.noreply.github.com>
Co-authored-by: Matt Robinson <mrobinson@unstructuredai.io>
Co-authored-by: Jeff Huber <jeffchuber@gmail.com>
Co-authored-by: Akshay <64036106+akshayvkt@users.noreply.github.com>
Co-authored-by: Andrew Huang <jhuang16888@gmail.com>
Co-authored-by: rogerserper <124558887+rogerserper@users.noreply.github.com>
Co-authored-by: seanaedmiston <seane999@gmail.com>
Co-authored-by: Hasegawa Yuya <52068175+Hase-U@users.noreply.github.com>
Co-authored-by: Ivan Vendrov <ivendrov@gmail.com>
Co-authored-by: Chen Wu (吴尘) <henrychenwu@cmu.edu>
Co-authored-by: Dennis Antela Martinez <dennis.antela@gmail.com>
Co-authored-by: Maxime Vidal <max.vidal@hotmail.fr>
Co-authored-by: Rishabh Raizada <110235735+rishabh-ti@users.noreply.github.com>

tests/integration_tests/embeddings/test_self_hosted.py

@@ -0,0 +1,96 @@
+"""Test self-hosted embeddings."""
+from typing import Any
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+from langchain.embeddings import (
+    SelfHostedEmbeddings,
+    SelfHostedHuggingFaceEmbeddings,
+    SelfHostedHuggingFaceInstructEmbeddings,
+)
+
+
+def get_remote_instance() -> Any:
+    """Get remote instance for testing."""
+    import runhouse as rh
+
+    gpu = rh.cluster(name="rh-a10x", instance_type="A100:1", use_spot=False)
+    gpu.install_packages(["pip:./"])
+    return gpu
+
+
+def test_self_hosted_huggingface_embedding_documents() -> None:
+    """Test self-hosted huggingface embeddings."""
+    documents = ["foo bar"]
+    gpu = get_remote_instance()
+    embedding = SelfHostedHuggingFaceEmbeddings(hardware=gpu)
+    output = embedding.embed_documents(documents)
+    assert len(output) == 1
+    assert len(output[0]) == 768
+
+
+def test_self_hosted_huggingface_embedding_query() -> None:
+    """Test self-hosted huggingface embeddings."""
+    document = "foo bar"
+    gpu = get_remote_instance()
+    embedding = SelfHostedHuggingFaceEmbeddings(hardware=gpu)
+    output = embedding.embed_query(document)
+    assert len(output) == 768
+
+
+def test_self_hosted_huggingface_instructor_embedding_documents() -> None:
+    """Test self-hosted huggingface instruct embeddings."""
+    documents = ["foo bar"]
+    gpu = get_remote_instance()
+    embedding = SelfHostedHuggingFaceInstructEmbeddings(hardware=gpu)
+    output = embedding.embed_documents(documents)
+    assert len(output) == 1
+    assert len(output[0]) == 768
+
+
+def test_self_hosted_huggingface_instructor_embedding_query() -> None:
+    """Test self-hosted huggingface instruct embeddings."""
+    query = "foo bar"
+    gpu = get_remote_instance()
+    embedding = SelfHostedHuggingFaceInstructEmbeddings(hardware=gpu)
+    output = embedding.embed_query(query)
+    assert len(output) == 768
+
+
+def get_pipeline() -> Any:
+    """Get pipeline for testing."""
+    model_id = "facebook/bart-base"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+    return pipeline("feature-extraction", model=model, tokenizer=tokenizer)
+
+
+def inference_fn(pipeline: Any, prompt: str) -> Any:
+    """Inference function for testing."""
+    # Return last hidden state of the model
+    if isinstance(prompt, list):
+        return [emb[0][-1] for emb in pipeline(prompt)]
+    return pipeline(prompt)[0][-1]
+
+
+def test_self_hosted_embedding_documents() -> None:
+    """Test self-hosted huggingface instruct embeddings."""
+    documents = ["foo bar"] * 2
+    gpu = get_remote_instance()
+    embedding = SelfHostedEmbeddings(
+        model_load_fn=get_pipeline, hardware=gpu, inference_fn=inference_fn
+    )
+    output = embedding.embed_documents(documents)
+    assert len(output) == 2
+    assert len(output[0]) == 50265
+
+
+def test_self_hosted_embedding_query() -> None:
+    """Test self-hosted custom embeddings."""
+    query = "foo bar"
+    gpu = get_remote_instance()
+    embedding = SelfHostedEmbeddings(
+        model_load_fn=get_pipeline, hardware=gpu, inference_fn=inference_fn
+    )
+    output = embedding.embed_query(query)
+    assert len(output) == 50265

tests/integration_tests/llms/test_self_hosted_llm.py

@@ -0,0 +1,105 @@
+"""Test Self-hosted LLMs."""
+import pickle
+from typing import Any, List, Optional
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+from langchain.llms import SelfHostedHuggingFaceLLM, SelfHostedPipeline
+
+model_reqs = ["pip:./", "transformers", "torch"]
+
+
+def get_remote_instance() -> Any:
+    """Get remote instance for testing."""
+    import runhouse as rh
+
+    return rh.cluster(name="rh-a10x", instance_type="A100:1", use_spot=False)
+
+
+def test_self_hosted_huggingface_pipeline_text_generation() -> None:
+    """Test valid call to self-hosted HuggingFace text generation model."""
+    gpu = get_remote_instance()
+    llm = SelfHostedHuggingFaceLLM(
+        model_id="gpt2",
+        task="text-generation",
+        model_kwargs={"n_positions": 1024},
+        hardware=gpu,
+        model_reqs=model_reqs,
+    )
+    output = llm("Say foo:")  # type: ignore
+    assert isinstance(output, str)
+
+
+def test_self_hosted_huggingface_pipeline_text2text_generation() -> None:
+    """Test valid call to self-hosted HuggingFace text2text generation model."""
+    gpu = get_remote_instance()
+    llm = SelfHostedHuggingFaceLLM(
+        model_id="google/flan-t5-small",
+        task="text2text-generation",
+        hardware=gpu,
+        model_reqs=model_reqs,
+    )
+    output = llm("Say foo:")  # type: ignore
+    assert isinstance(output, str)
+
+
+def load_pipeline() -> Any:
+    """Load pipeline for testing."""
+    model_id = "gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+    pipe = pipeline(
+        "text-generation", model=model, tokenizer=tokenizer, max_new_tokens=10
+    )
+    return pipe
+
+
+def inference_fn(pipeline: Any, prompt: str, stop: Optional[List[str]] = None) -> str:
+    """Inference function for testing."""
+    return pipeline(prompt)[0]["generated_text"]
+
+
+def test_init_with_local_pipeline() -> None:
+    """Test initialization with a self-hosted HF pipeline."""
+    gpu = get_remote_instance()
+    pipeline = load_pipeline()
+    llm = SelfHostedPipeline.from_pipeline(
+        pipeline=pipeline,
+        hardware=gpu,
+        model_reqs=model_reqs,
+        inference_fn=inference_fn,
+    )
+    output = llm("Say foo:")  # type: ignore
+    assert isinstance(output, str)
+
+
+def test_init_with_pipeline_path() -> None:
+    """Test initialization with a self-hosted HF pipeline."""
+    gpu = get_remote_instance()
+    pipeline = load_pipeline()
+    import runhouse as rh
+
+    rh.blob(pickle.dumps(pipeline), path="models/pipeline.pkl").save().to(
+        gpu, path="models"
+    )
+    llm = SelfHostedPipeline.from_pipeline(
+        pipeline="models/pipeline.pkl",
+        hardware=gpu,
+        model_reqs=model_reqs,
+        inference_fn=inference_fn,
+    )
+    output = llm("Say foo:")  # type: ignore
+    assert isinstance(output, str)
+
+
+def test_init_with_pipeline_fn() -> None:
+    """Test initialization with a self-hosted HF pipeline."""
+    gpu = get_remote_instance()
+    llm = SelfHostedPipeline(
+        model_load_fn=load_pipeline,
+        hardware=gpu,
+        model_reqs=model_reqs,
+        inference_fn=inference_fn,
+    )
+    output = llm("Say foo:")  # type: ignore
+    assert isinstance(output, str)