(
lm_config: lm_config.LMConfig,
prompt: APIInput,
)
| 16 | tokenizer = None |
| 17 | |
| 18 | def call_llm( |
| 19 | lm_config: lm_config.LMConfig, |
| 20 | prompt: APIInput, |
| 21 | ) -> str: |
| 22 | global model |
| 23 | global tokenizer |
| 24 | |
| 25 | response: str |
| 26 | |
| 27 | if lm_config.provider == "openai": |
| 28 | if lm_config.mode == "chat": |
| 29 | assert isinstance(prompt, list) |
| 30 | response = generate_from_openai_chat_completion( |
| 31 | messages=prompt, |
| 32 | model=lm_config.model, |
| 33 | temperature=lm_config.gen_config["temperature"], |
| 34 | top_p=lm_config.gen_config["top_p"], |
| 35 | context_length=lm_config.gen_config["context_length"], |
| 36 | max_tokens=lm_config.gen_config["max_tokens"], |
| 37 | stop_token=None, |
| 38 | ) |
| 39 | elif lm_config.mode == "completion": |
| 40 | assert isinstance(prompt, str) |
| 41 | response = generate_from_openai_completion( |
| 42 | prompt=prompt, |
| 43 | engine=lm_config.model, |
| 44 | temperature=lm_config.gen_config["temperature"], |
| 45 | max_tokens=lm_config.gen_config["max_tokens"], |
| 46 | top_p=lm_config.gen_config["top_p"], |
| 47 | stop_token=lm_config.gen_config["stop_token"], |
| 48 | ) |
| 49 | else: |
| 50 | raise ValueError( |
| 51 | f"OpenAI models do not support mode {lm_config.mode}" |
| 52 | ) |
| 53 | elif lm_config.provider == "huggingface": |
| 54 | assert isinstance(prompt, str) |
| 55 | response = generate_from_huggingface_completion( |
| 56 | prompt=prompt, |
| 57 | model_endpoint=lm_config.gen_config["model_endpoint"], |
| 58 | temperature=lm_config.gen_config["temperature"], |
| 59 | top_p=lm_config.gen_config["top_p"], |
| 60 | stop_sequences=lm_config.gen_config["stop_sequences"], |
| 61 | max_new_tokens=lm_config.gen_config["max_new_tokens"], |
| 62 | ) |
| 63 | elif lm_config.provider == "ours": |
| 64 | # print(prompt) |
| 65 | if lm_config.model == 'manual': |
| 66 | response = input("Command > ") |
| 67 | else: |
| 68 | if not model: |
| 69 | model = AutoModel.from_pretrained(lm_config.model, trust_remote_code=True, device=f'cuda:{lm_config.cuda}') |
| 70 | tokenizer = AutoTokenizer.from_pretrained(lm_config.model, trust_remote_code=True) |
| 71 | model.eval() |
| 72 | response = call_pretrain_model(prompt, model, tokenizer, lm_config.cuda) |
| 73 | else: |
| 74 | raise NotImplementedError( |
| 75 | f"Provider {lm_config.provider} not implemented" |
no test coverage detected