HFChatbot.py

from ragpipeline import ChatBot
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

class HFBot(ChatBot):
    def __init__(self, 
                 model_name: str = "Almawave/Velvet-2B",
                 model_system=(
                "Sei un esperto di diritto amministrativo che deve eseguire il "
                "controllo di regolarità amministrativa su un atto amministrativo di un comune italiano. "
                "Ti verranno forniti un atto amministrativo  (determinazione dirigenziale) ed eventuali allegati, questi sono forniti come frammenti rilevanti. "
                "Utilizza solamente i frammenti che ti verranno inviati."
                "Rispondi in Italiano usando al massimo 50 parole. "
                "Basati esclusivamente sul seguente testo: "
            ),
                dump_filename="dump.csv"
                     ):
        self.model_name = model_name
        self.model_orig = model_name
        self.model_system = model_system
        self.dump_filename = dump_filename
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
       
    def chat(self, domanda: str, istruzioni: str = None, frammenti =[]) -> str:
        prompt = f"ISTRUZIONI: {istruzioni}\n\nCONTESTO:\n" + "\n".join(frammenti) + f"\n\nDOMANDA: {domanda}"
        
        response = self.genera(prompt)
        
        return response

    def generate(self,
             relevant_docs = [],
             attributi_frammenti_rilevanti = [],
             query="",
             istruzioni :str = None ##togliere
             ):
        
        i = 0
        #print (f"DIMESIONE FILE {len(relevant_files)}")
        #print (f"DIMESIONE TESTI {len(relevant_docs)}")
        
        prompt=""
        for documento in relevant_docs:
            prompt += f"{relevant_docs[i]} "
            i = i+1
        #"{context}\n\nDomanda: {query}"
        if istruzioni is not None:
            query = query +  " Istruzioni: " + istruzioni
        prompt +=f"\n\nDomanda:{query} \n\n"
        #print(prompt)
        rersponse = self.genera(prompt)
        return rersponse
    def genera(self, prompt: str):  
        retval = ""
        messages = [
            {"role": "system", "content": self.model_system,},
            {"role": "user", "content": prompt},
        ]
        


        pipe = pipeline("text-generation", model=self.model_name)

        output = pipe(messages)
        output
        for item in output:
            for entry in item['generated_text']:
                if entry['role'] == 'assistant':
                    retval = entry['content']
        return retval