test: auto gptq

utils/epfl_meditron_utils.py

@@ -1,38 +1,47 @@
-from ctransformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import pipeline
-import streamlit as st
-from langchain.chains import LLMChain
-from langchain.prompts import PromptTemplate
-
-# Simple inference example
-# output = llm(
-#  "<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant", # Prompt
-#  max_tokens=512,  # Generate up to 512 tokens
-#  stop=["</s>"],   # Example stop token - not necessarily correct for this specific model! Please check before using.
-#  echo=True        # Whether to echo the prompt
-#)
-
-prompt_format = "<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant"
-
-
-template = """Question: {question}
-
-Answer:"""
-
 
 
 
 def get_llm_response(repo, filename, model_type, gpu_layers, prompt):
-    print("Loading model")
-    llm = AutoModelForCausalLM.from_pretrained(repo, model_file=filename, model_type=model_type, gpu_layers=gpu_layers)
-    print("Model loaded")
-
-    #llm_prompt = prompt_format.format(system_message=system_prompt, prompt=prompt)
-    print(f"LLM prompt: {prompt}")
-    
-    prompt = PromptTemplate(template=template, input_variables=["question"])
-
-    llm_chain = LLMChain(prompt=prompt, llm=llm)
-    response = llm_chain.run(prompt)
-    
-    return response
+    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+    model_name_or_path = "TheBloke/meditron-7B-GPTQ"
+    # To use a different branch, change revision
+    # For example: revision="gptq-4bit-128g-actorder_True"
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
+                                                device_map="auto",
+                                                trust_remote_code=False,
+                                                revision="main")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)    
+
+    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
+
+    print("*** Pipeline:")
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+        top_k=40,
+        repetition_penalty=1.1
+    )
+
+    prompt_template=f'''<|im_start|>system
+    {system_message}<|im_end|>
+    <|im_start|>user
+    {prompt}<|im_end|>
+    <|im_start|>assistant
+    '''.format(system_message="You are an assistant", prompt=prompt)
+
+    response = pipe(prompt_template)[0]['generated_text']
+    print(response)
+    return response