Spaces:
Runtime error
Runtime error
File size: 2,363 Bytes
96b7a63 a7cdd19 c9fffcd a7cdd19 c9fffcd 96b7a63 |
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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from typing import Any, List, Mapping, Optional
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from typing import Literal
import requests
from langchain.prompts import PromptTemplate, ChatPromptTemplate
from operator import itemgetter
import re, os
def format_captions(text):
sentences = list(filter(None,[x.strip() for x in re.split(r'[^A-Za-z0-9 -]', text)]))
print(len(sentences))
return "\n".join([f"{i+1}. {d}" for i,d in enumerate(sentences)])
def custom_chain():
API_TOKEN = os.environ['HF_INFER_API']
prompt = PromptTemplate.from_template("<s><INST>Given the below template, create a list of image generation prompt with maximum 5 words for each number\n\n{template}<INST> ")
cap_llm = CustomLLM(repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1", model_type='text-generation', api_token=API_TOKEN, stop=["\n<|"])
return {"template":lambda x:format_captions(x)} | prompt | cap_llm
class CustomLLM(LLM):
repo_id : str
api_token : str
model_type: Literal["text2text-generation", "text-generation"]
max_new_tokens: int = None
temperature: float = 0.001
timeout: float = None
top_p: float = None
top_k : int = None
repetition_penalty : float = None
stop : List[str] = []
@property
def _llm_type(self) -> str:
return "custom"
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
headers = {"Authorization": f"Bearer {self.api_token}"}
API_URL = f"https://api-inference.huggingface.co/models/{self.repo_id}"
parameters_dict = {
'max_new_tokens': self.max_new_tokens,
'temperature': self.temperature,
'timeout': self.timeout,
'top_p': self.top_p,
'top_k': self.top_k,
'repetition_penalty': self.repetition_penalty,
'stop':self.stop
}
if self.model_type == 'text-generation':
parameters_dict["return_full_text"]=False
data = {"inputs": prompt, "parameters":parameters_dict, "options":{"wait_for_model":True}}
data = requests.post(API_URL, headers=headers, json=data).json()
return data[0]['generated_text'] |