""" title: Huggingface Pipe author: Wes Caldwell email: Musicheardworldwide@gmail.com 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)