"""
title: Huggingface Pipe
author: Wes Caldwell
email:
[email protected]
date: 2024-07-19
version: 1.05
license: MIT
description: Function to use HF models on Open-WebUI
requirements: pydantic, huggingface-hub, transformers
"""
from pydantic import BaseModel, Field
from typing import Optional, Union, Generator, Iterator, List
import os
import logging
from huggingface_hub import HfApi
from transformers import pipeline, set_seed, PipelineException
logging.basicConfig(level=logging.INFO)
class Pipe:
class Valves(BaseModel):
NAME_PREFIX: str = Field(
default="HUGGINGFACE/",
description="Prefix to be added before model names.",
)
HUGGINGFACE_API_URL: str = Field(
default="https://api-inference.huggingface.co/models/",
description="Base URL for accessing Hugging Face API endpoints.",
)
HUGGINGFACE_API_KEY: str = Field(
default=os.getenv("HUGGINGFACE_API_KEY", ""),
description="API key for authenticating requests to the Hugging Face API.",
)
def __init__(self):
self.type = "manifold"
self.valves = self.Valves()
self.hf_api = HfApi()
def fetch_models(self) -> List[dict]:
"""Fetch models from Hugging Face containing 'gpt' in their ID."""
if not self.valves.HUGGINGFACE_API_KEY:
logging.error("API Key not provided.")
return [
{"id": "error", "name": "API Key not provided."}
]
try:
models = self.hf_api.list_models(
use_auth_token=self.valves.HUGGINGFACE_API_KEY
)
filtered_models = [
{"id": model.modelId, "name": f"{self.valves.NAME_PREFIX}{model.modelId}"}
for model in models
if "gpt" in model.modelId
]
return filtered_models
except Exception as e:
logging.error(f"Failed to fetch models: {e}")
return [
{"id": "error", "name": "Could not fetch models. Update the API Key."}
]
def pipe(self, body: dict, __user__: dict) -> Union[str, Generator, Iterator]:
"""Process a text-generation request."""
if "model" not in body:
logging.error("Model not specified in the request body.")
return "Error: Model not specified in the request body."
if "prompt" not in body:
logging.error("Prompt not provided in the request body.")
return "Error: Prompt not provided in the request body."
model_id = body["model"]
logging.info(f"Processing request for model: {model_id}")
try:
# Set up the text generation pipeline
device = 0 if os.getenv("CUDA_VISIBLE_DEVICES") else -1
generator = pipeline(
"text-generation",
model=model_id,
tokenizer=model_id,
device=device,
)
set_seed(42) # Optional: Ensure reproducibility
response = generator(
body["prompt"],
max_length=body.get("max_tokens", 50),
num_return_sequences=1,
do_sample=True,
)
if body.get("stream", False):
# Stream results as a generator
for message in response:
yield message["generated_text"]
else:
# Return the entire response
return [msg["generated_text"] for msg in response]
except PipelineException as pe:
logging.error(f"Pipeline error: {pe}")
return f"Error: {pe}"
except Exception as e:
logging.error(f"Failed to process request: {e}")
return f"Error: {e}"
if __name__ == "__main__":
pipe = Pipe()
# Test fetch_models
models = pipe.fetch_models()
print(models)
# Test pipe method
test_body = {
"model": "gpt2",
"prompt": "Once upon a time",
"max_tokens": 50,
"stream": False,
}
result = pipe.pipe(test_body, __user__={})
if isinstance(result, str):
print(result)
else:
for res in result:
print(res)
