File size: 2,039 Bytes
65d97fa 3afd61a 65d97fa 3afd61a 65d97fa 590d088 65d97fa f4f9c98 65d97fa |
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 |
import os
from .amodel import AModel
from huggingface_hub import InferenceClient
import numpy as np # feature_extraction renvoie un array numpy...
class HuggingFaceModel(AModel):
def __init__(self, llm_name:str, feature_name:str, temperature:float=0.0):
self.llm_name:str = llm_name
self.feature_name:str = feature_name
self.temperature = temperature
api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
try:
self.model = InferenceClient(api_key=api_token)
except:
raise
def ask_llm(self, question:str)->str:
messages = [{"role": "user", "content": question}]
try:
resp = self.model.chat.completions.create(
model=self.llm_name,
messages=messages,
max_tokens=2048,
temperature=self.temperature,
# stream=True
)
return resp.choices[0].message.content
except:
raise
def create_vector(self, chunk:str)->list[float]:
try:
resp = self.model.feature_extraction(
text=chunk,
# normalize=True, # Only available on server powered by Text-Embedding-Inference.
model=self.feature_name, # normalisé ??
)
if isinstance(resp, np.ndarray):
return resp
else:
raise Exception("Error with embedding !")
except:
raise
def create_vectors(self, chunks:list[str])->list[list[float]]:
'''
Pas de batch pour la création de vectors sur HuggingFace, on les passe un par un
'''
vectors = []
try:
for chunk in chunks:
v = self.create_vector(chunk)
if not isinstance(v, np.ndarray): # l'api renvoie des array numpy....
raise
vectors.append(v.tolist())
return vectors
except:
raise |