Merge pull request #11 from luisrodriguesphd/refact-params-tg

conf/params.yml

@@ -16,12 +16,16 @@ embedding_dir: ["data", "03_indexed"]
 
 
 # LLM / Text Generation
-model_id: "M4-ai/tau-1.8B" 
 # See instructions for parameters: https://www.ibm.com/docs/en/watsonx-as-a-service?topic=lab-model-parameters-prompting
-top_k: 30 
-top_p: 0.7
-temperature: 0.3
-max_new_tokens: 256
+generative_model:
+    model_name: "M4-ai/tau-1.8B"
+    model_kwargs:
+        trust_remote_code: True
+    generate_kwargs: 
+        top_k: 30 
+        top_p: 0.7
+        temperature: 0.3
+        max_new_tokens: 256
 # See instructions for the prompt: https://huggingface.co/spaces/Locutusque/Locutusque-Models/blob/main/app.py
 prompt_dir: ["data", "04_prompts"]
 promp_file: "prompt_template_for_explaning_why_is_a_good_fit.json"

src/resume_worth/pipelines/text_generation/nodes.py

@@ -1,4 +1,6 @@
 import os
+os.environ['HF_HOME'] = ".cache/huggingface"
+
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain_core.prompts import PromptTemplate
@@ -10,20 +12,34 @@ import transformers
 transformers.logging.set_verbosity_error()
 
 
-@lru_cache(maxsize=None)
-def load_hf_text_generation_model_to_langchain(model_id:str='gpt2', top_k:int=50, top_p:float=0.95, temperature:float=0.4, max_new_tokens:int=1024):
+#@lru_cache(maxsize=None)
+def load_hf_text_generation_model_to_langchain(
+    model_name:str='gpt2', 
+    model_kwargs:dict={
+            'trust_remote_code': True,
+    },
+    generate_kwargs:dict={
+            'top_k': 50, 
+            'top_p': 0.95, 
+            'temperature': 0.4, 
+            'max_new_tokens': 1024,
+        }
+    ):
     """
     Function to load a text generation model hosted on Hugging Face to se used in LangChain.
     More info, see: https://python.langchain.com/docs/integrations/llms/huggingface_pipelines/
     """
 
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id)
+    print(f"-> Load a pretrained text embedding model {model_name}")
+
+    # https://huggingface.co/apple/OpenELM
+    tokenizer = AutoTokenizer.from_pretrained(model_name, **model_kwargs)
+    model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
     
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, 
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer,
             return_full_text=False, do_sample=True, 
-            top_p=top_p, top_k=top_k, temperature=temperature, max_new_tokens=max_new_tokens, 
-            num_beams=1, repetition_penalty=1.1, num_return_sequences=1
+            **generate_kwargs,
+            num_beams=1, repetition_penalty=1.1, num_return_sequences=1,
         )
     
     hf = HuggingFacePipeline(pipeline=pipe)

src/resume_worth/pipelines/text_generation/pipeline.py

@@ -11,11 +11,7 @@ from resume_worth.pipelines.text_generation.nodes import load_hf_text_generation
 
 
 params = get_params()
-model_id = params['model_id']
-top_p = params['top_p']
-top_k = params['top_k']
-temperature = params['temperature']
-max_new_tokens = params['max_new_tokens']
+generative_model = params['generative_model']
 prompt_dir = params['prompt_dir']
 promp_file = params['promp_file']
 
@@ -24,7 +20,7 @@ def generate_explanation_why_resume_for_a_job(resume: str, job: str):
 
     # Stage 1 - [cacheable] Load text generation model
 
-    text_generation_model = load_hf_text_generation_model_to_langchain(model_id, top_k, top_p, temperature, max_new_tokens)
+    text_generation_model = load_hf_text_generation_model_to_langchain(generative_model['model_name'], generative_model['model_kwargs'], generative_model['generate_kwargs'])
 
     # Stage 2 - [cacheable] Load text generation model
     

src/resume_worth/utils/utils.py

@@ -1,4 +1,6 @@
 import os
+os.environ['HF_HOME'] = ".cache/huggingface"
+
 import yaml
 from langchain_community.embeddings import HuggingFaceEmbeddings
 import fitz # imports the pymupdf library