You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
51 lines
1.6 KiB
51 lines
1.6 KiB
from pathlib import Path
|
|
|
|
import pytest
|
|
from models.openai_model import Model
|
|
from transformers import GPT2TokenizerFast
|
|
|
|
from services.usage_service import UsageService
|
|
|
|
|
|
# Non-ChatGPT -> TODO: make generic test and loop through text models
|
|
@pytest.mark.asyncio
|
|
async def test_send_req():
|
|
usage_service = UsageService(Path("../tests"))
|
|
model = Model(usage_service)
|
|
prompt = "how many hours are in a day?"
|
|
tokens = len(GPT2TokenizerFast.from_pretrained("gpt2")(prompt)["input_ids"])
|
|
res = await model.send_request(prompt, tokens)
|
|
assert "24" in res["choices"][0]["text"]
|
|
|
|
|
|
# ChatGPT version
|
|
@pytest.mark.asyncio
|
|
async def test_send_req_gpt():
|
|
usage_service = UsageService(Path("../tests"))
|
|
model = Model(usage_service)
|
|
prompt = "how many hours are in a day?"
|
|
res = await model.send_request(
|
|
prompt, None, is_chatgpt_request=True, model="gpt-3.5-turbo"
|
|
)
|
|
assert "24" in res["choices"][0]["message"]["content"]
|
|
|
|
|
|
# GPT4 version
|
|
@pytest.mark.asyncio
|
|
async def test_send_req_gpt4():
|
|
usage_service = UsageService(Path("../tests"))
|
|
model = Model(usage_service)
|
|
prompt = "how many hours are in a day?"
|
|
res = await model.send_request(prompt, None, is_chatgpt_request=True, model="gpt-4")
|
|
assert "24" in res["choices"][0]["message"]["content"]
|
|
|
|
|
|
# Edit request -> currently broken due to endpoint
|
|
# @pytest.mark.asyncio
|
|
# async def test_send_edit_req():
|
|
# usage_service = UsageService(Path("../tests"))
|
|
# model = Model(usage_service)
|
|
# text = 'how many hours are in a day?'
|
|
# res = await model.send_edit_request(text)
|
|
# assert '24' in res['choices'][0]['text']
|