58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
# coding=utf-8
|
|
# Copyright 2024 HuggingFace Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import json
|
|
import unittest
|
|
from typing import Optional
|
|
|
|
from smolagents import ChatMessage, HfApiModel, TransformersModel, models, tool
|
|
|
|
|
|
class ModelTests(unittest.TestCase):
|
|
def test_get_json_schema_has_nullable_args(self):
|
|
@tool
|
|
def get_weather(location: str, celsius: Optional[bool] = False) -> str:
|
|
"""
|
|
Get weather in the next days at given location.
|
|
Secretly this tool does not care about the location, it hates the weather everywhere.
|
|
|
|
Args:
|
|
location: the location
|
|
celsius: the temperature type
|
|
"""
|
|
return "The weather is UNGODLY with torrential rains and temperatures below -10°C"
|
|
|
|
assert "nullable" in models.get_json_schema(get_weather)["function"]["parameters"]["properties"]["celsius"]
|
|
|
|
def test_chatmessage_has_model_dumps_json(self):
|
|
message = ChatMessage("user", "Hello!")
|
|
data = json.loads(message.model_dump_json())
|
|
assert data["content"] == "Hello!"
|
|
|
|
def test_get_hfapi_message_no_tool(self):
|
|
model = HfApiModel(max_tokens=10)
|
|
messages = [{"role": "user", "content": "Hello!"}]
|
|
model(messages, stop_sequences=["great"])
|
|
|
|
def test_transformers_message_no_tool(self):
|
|
model = TransformersModel(
|
|
model_id="HuggingFaceTB/SmolLM2-135M-Instruct",
|
|
max_new_tokens=5,
|
|
device_map="auto",
|
|
do_sample=False,
|
|
)
|
|
messages = [{"role": "user", "content": "Hello!"}]
|
|
output = model(messages, stop_sequences=["great"]).content
|
|
assert output == "assistant\nHello"
|