Spaces:
Sleeping
Sleeping
File size: 772 Bytes
0809507 ef2a3f4 0809507 ef2a3f4 0809507 ef2a3f4 0809507 |
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 |
from typing import Any, List, Mapping, Optional
from langchain.llms.base import LLM
from transformers import Pipeline
class CustomLLM(LLM):
pipeline = None
# Create the pipeline for question answering
def __init__(self, model_pipeline: Pipeline):
self.pipeline = model_pipeline
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
prompt_length = len(prompt)
response = self.pipeline(prompt, max_new_tokens=525)[0]["generated_text"]
# only return newly generated tokens
return response[prompt_length:]
@property
def _identifying_params(self) -> Mapping[str, Any]:
return {"name_of_model": self.model_name}
@property
def _llm_type(self) -> str:
return "custom" |