from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

class HFBot():
    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: "
            ),
                
                     ):
        self.model_name = model_name
        self.model_system = model_system
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
       
    def chat(self, domanda: str, istruzioni: str = None, frammenti =[]) -> str:
        prompt = f"\nISTRUZIONI: {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="\nCONTESTO: "
        for documento in relevant_docs:
            prompt += f"{relevant_docs[i]} "
            i = i+1
        #"{context}\n\nDomanda: {query}"
        if istruzioni is not None:
            query = query +  " \nISTRUZIONI: " + istruzioni
        prompt +=f"\n\DOMANDA:{query} \n\n"
        #print(prompt)
        rersponse = self.genera(prompt)
        return rersponse
    
    def riassumi(self, text):
        retval =""

        pipe = pipeline("summarization", model=self.model_name)
        retval = pipe(text)
        return retval
    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