Spaces:
Sleeping
Sleeping
File size: 2,062 Bytes
0d17078 beaf47c 0d17078 cdd1ac6 beaf47c 0d17078 cdd1ac6 0d17078 beaf47c cdd1ac6 0d17078 cdd1ac6 0d17078 beaf47c 0d17078 beaf47c 0d17078 cdd1ac6 beaf47c 0d17078 cdd1ac6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun
from typing import Any, List, Optional, Dict
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from pydantic import Field, PrivateAttr
class HuggingFaceLLM(LLM):
model_id: str = Field(..., description="Hugging Face model ID")
temperature: float = Field(default=0.7, description="Sampling temperature")
max_tokens: int = Field(default=256, description="Maximum number of tokens to generate")
device: str = Field(default="cpu", description="Device to run the model on")
_model: Any = PrivateAttr()
_tokenizer: Any = PrivateAttr()
def __init__(self, **kwargs):
super().__init__(**kwargs)
if self.device == "cpu":
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self._load_model()
def _load_model(self):
self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
self._model = AutoModelForCausalLM.from_pretrained(self.model_id).to(self.device)
@property
def _llm_type(self) -> str:
return "custom_huggingface"
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
input_ids = self._tokenizer.encode(prompt, return_tensors="pt").to(self.device)
with torch.no_grad():
output = self._model.generate(
input_ids,
max_new_tokens=self.max_tokens,
temperature=self.temperature,
do_sample=True,
pad_token_id=self._tokenizer.eos_token_id
)
response = self._tokenizer.decode(output[0], skip_special_tokens=True)
return response[len(prompt):].strip()
@property
def _identifying_params(self) -> Dict[str, Any]:
return {"model_id": self.model_id, "temperature": self.temperature, "max_tokens": self.max_tokens, "device": self.device} |