|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import sys, os |
|
from pprint import pprint |
|
from jinja2 import Environment, FileSystemLoader, meta |
|
import yaml |
|
|
|
import pandas as pd |
|
from vllm import LLM, SamplingParams |
|
|
|
|
|
sys.path.append(".") |
|
os.chdir(os.path.dirname(os.path.abspath(__file__))) |
|
|
|
|
|
def get_llm_response(prompt_template, sampling_params): |
|
prompts = [prompt_template] |
|
outputs = llm.generate(prompts, sampling_params, use_tqdm = False) |
|
generated_text = outputs[0].outputs[0].text |
|
if generated_text[-1] != "}": |
|
generated_text = generated_text + "}" |
|
prompt = prompt_template + generated_text |
|
return prompt, generated_text |
|
|
|
if __name__ == "__main__": |
|
|
|
with open('prompt_config.yaml') as f: |
|
config = yaml.safe_load(f) |
|
|
|
print("prompt format:", config.get("prompt_format")) |
|
print(config) |
|
print() |
|
for prompt in config["prompts"]: |
|
if prompt["mode"] == "analysis": |
|
print(f'--- prompt mode: {prompt["mode"]} ---') |
|
env = Environment(loader=FileSystemLoader(".")) |
|
template = env.get_template(prompt["template"]) |
|
|
|
source = template.environment.loader.get_source(template.environment, template.name) |
|
variables = meta.find_undeclared_variables(env.parse(source[0])) |
|
|
|
print("variables:", variables) |
|
print("---") |
|
|
|
data = {"query": "Comment obtenir le formulaire A36 ?"} |
|
if "system_prompt" in variables: |
|
data["system_prompt"] = prompt["system_prompt"] |
|
|
|
rendered_template = template.render(**data) |
|
print(rendered_template) |
|
print("---") |
|
|
|
llm = LLM("mistral-mfs-reference-2/mistral-mfs-reference-2") |
|
|
|
sampling_params = SamplingParams(temperature=0.2, top_p=0.95, max_tokens=300, stop="}") |
|
|
|
prompt, generated_text = get_llm_response(rendered_template, sampling_params) |
|
print("Albert : ", generated_text) |
|
|
