test: generation

utils/epfl_meditron_utils.py

@@ -1,8 +1,7 @@
-
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 
 def get_llm_response(repo, filename, model_type, gpu_layers, system_message, prompt):
-    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
     model_name_or_path = "TheBloke/meditron-7B-GPTQ"
     # To use a different branch, change revision
@@ -12,15 +11,23 @@ def get_llm_response(repo, filename, model_type, gpu_layers, system_message, pro
                                                 trust_remote_code=False,
                                                 revision="main")
 
-    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)    
-
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
+    
+    prompt_template=f'''<|im_start|>system
+        {system_message}<|im_end|>
+        <|im_start|>user
+        {prompt}<|im_end|>
+        <|im_start|>assistant
+        '''
+        
+    print("Template:")
+    print(prompt_template)
+    
     print("\n\n*** Generate:")
 
-    #input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
-    #output = model.generate(inputs=input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
-    #print(tokenizer.decode(output[0]))
-
-    # Inference can also be done using transformers' pipeline
+    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
+    output = model.generate(inputs=input_ids, temperature=0.01, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
+    print(tokenizer.decode(output[0])) 
 
     print("*** Pipeline:")
     pipe = pipeline(
@@ -35,12 +42,7 @@ def get_llm_response(repo, filename, model_type, gpu_layers, system_message, pro
         repetition_penalty=1.1
     )
 
-    prompt_template=f'''<|im_start|>system
-    {system_message}<|im_end|>
-    <|im_start|>user
-    {prompt}<|im_end|>
-    <|im_start|>assistant
-    '''
+    
 
     response = pipe(prompt_template)[0]['generated_text']
     print(response)