src/model_huggingface.py

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