diff --git a/llmx/generators/text/ollama_textgen.py b/llmx/generators/text/ollama_textgen.py new file mode 100644 index 0000000..8744f19 --- /dev/null +++ b/llmx/generators/text/ollama_textgen.py @@ -0,0 +1,85 @@ +import os +import requests +from dataclasses import asdict +from typing import Union, List, Dict +from .base_textgen import TextGenerator +from ...datamodel import Message, TextGenerationConfig, TextGenerationResponse +from ...utils import cache_request, get_models_maxtoken_dict, num_tokens_from_messages + + +class OllamaTextGenerator(TextGenerator): + def __init__( + self, + model: str = None, + base_url: str = None, + models: Dict = None, + **kwargs + ): + super().__init__(provider="ollama") + self.model_name = model or "gemma2" + self.base_url = base_url or os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434") + self.model_max_token_dict = get_models_maxtoken_dict(models) + + def generate( + self, + messages: Union[List[dict], str], + config: TextGenerationConfig = TextGenerationConfig(), + **kwargs, + ) -> TextGenerationResponse: + use_cache = config.use_cache + model = config.model or self.model_name + + prompt_tokens = num_tokens_from_messages(messages) + + max_tokens = max( + self.model_max_token_dict.get(model, 4096) - prompt_tokens - 10, 200 + ) + + ollama_payload = { + "model": model, + "prompt": messages if isinstance(messages, str) else messages[-1]["content"], # Assume last message + "temperature": config.temperature, + "max_tokens": max_tokens, + "top_p": config.top_p, + "stream": False, + } + + self.model_name = model + cache_key_params = (ollama_payload) | {"messages": messages} + + if use_cache: + response = cache_request(cache=self.cache, params=cache_key_params) + if response: + return TextGenerationResponse(**response) + + url = f"{self.base_url}/api/generate" + response = requests.post(url, json=ollama_payload) + oai_response = response.json() + + if response.status_code != 200: + raise ValueError(f"Ollama API error: {response.status_code}, {response.text}") + + total_tokens = oai_response.get("eval_count", 0) + prompt_tokens = oai_response.get("prompt_eval_count", 0) + + usage = { + "prompt_tokens": prompt_tokens, + "completion_tokens": total_tokens - prompt_tokens, + "total_tokens": total_tokens, + } + + response_obj = TextGenerationResponse( + text=[Message(content= oai_response["response"], role="assistant")], + logprobs=[], + config=ollama_payload, + usage=usage, + ) + + cache_request( + cache=self.cache, params=cache_key_params, values=asdict(response_obj) + ) + + return response_obj + + def count_tokens(self, text) -> int: + return num_tokens_from_messages(text) \ No newline at end of file diff --git a/llmx/generators/text/textgen.py b/llmx/generators/text/textgen.py index 3d86002..01d2b29 100644 --- a/llmx/generators/text/textgen.py +++ b/llmx/generators/text/textgen.py @@ -3,6 +3,7 @@ from .palm_textgen import PalmTextGenerator from .cohere_textgen import CohereTextGenerator from .anthropic_textgen import AnthropicTextGenerator +from .ollama_textgen import OllamaTextGenerator import logging logger = logging.getLogger("llmx") @@ -19,6 +20,8 @@ def sanitize_provider(provider: str): return "hf" elif provider.lower() == "anthropic" or provider.lower() == "claude": return "anthropic" + elif provider.lower() == "ollama": + return "ollama" else: raise ValueError( f"Invalid provider '{provider}'. Supported providers are 'openai', 'hf', 'palm', 'cohere', and 'anthropic'." @@ -58,6 +61,8 @@ def llm(provider: str = None, **kwargs): return CohereTextGenerator(**kwargs) elif provider.lower() == "anthropic": return AnthropicTextGenerator(**kwargs) + elif provider.lower() == "ollama": + return OllamaTextGenerator(**kwargs) elif provider.lower() == "hf": try: import transformers @@ -80,5 +85,5 @@ def llm(provider: str = None, **kwargs): else: raise ValueError( - f"Invalid provider '{provider}'. Supported providers are 'openai', 'hf', 'palm', 'cohere', and 'anthropic'." + f"Invalid provider '{provider}'. Supported providers are 'openai', 'hf', 'palm', 'cohere', 'anthropic', and 'ollama'." ) \ No newline at end of file diff --git a/notebooks/tutorial.ipynb b/notebooks/tutorial.ipynb index 82ac411..f6bf47b 100644 --- a/notebooks/tutorial.ipynb +++ b/notebooks/tutorial.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -12,7 +12,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 3, "metadata": {}, "outputs": [], "source": [ @@ -43,17 +43,9 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Gravity is like a big invisible force that pulls things towards each other. It's what keeps us on the ground and makes things fall down instead of floating away. Imagine if you threw a ball up in the air, gravity would pull it back down to the ground. It's like a super strong magnet that pulls everything together.\n" - ] - } - ], + "outputs": [], "source": [ "openai_gen = llm(provider=\"openai\", api_key=os.environ[\"OPENAI_API_KEY\"])\n", "openai_config = TextGenerationConfig(model=\"gpt-3.5-turbo\", use_cache=True)\n", @@ -70,17 +62,9 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Gravity is like a big invisible force that pulls things towards each other. It's what keeps us on the ground and makes things fall down when we drop them. It's like a big hug from the Earth that keeps us close to it.\n" - ] - } - ], + "outputs": [], "source": [ "from dotenv import load_dotenv \n", "\n", @@ -118,27 +102,9 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Gravity is a force that pulls objects towards each other. The more massive an object is, the stronger its gravitational pull. The Earth is very massive, so it has a strong gravitational pull. This is why we don't float off into space. The Moon is also massive, but it is much smaller than the Earth. This means that its gravitational pull is not as strong. This is why the Moon orbits the Earth, instead of the other way around.\n", - "\n", - "Gravity is a very important force in the universe. It is what keeps the planets in orbit around the Sun, and it is what keeps the Moon in orbit around the Earth. It is also what keeps us on the ground. Without gravity, we would all float off into space.\n", - "\n", - "Gravity is a very mysterious force. We don't really know what causes it. We do know that it is related to mass, but we don't know exactly how. Scientists are still working on trying to understand gravity.\n", - "\n", - "One way to think about gravity is to imagine a trampoline. If you put a bowling ball in the middle of the trampoline, it will make a dent in the trampoline. If you then put a marble on the trampoline, the marble will roll towards the bowling ball. This is because the bowling ball is more massive than the marble, and it has a stronger gravitational pull.\n", - "\n", - "The Earth is like the bowling ball, and we are like the marble. The Earth's gravity pulls us towards the center of the Earth, which is why we don't float off into space.\n", - "\n", - "Gravity is a very important force in the universe. It is what keeps the planets in orbit around the Sun, and it is what keeps the Moon in orbit around the Earth. It is also what keeps us on the ground. Without gravity, we would all float off into space.\n" - ] - } - ], + "outputs": [], "source": [ "palm_gen = llm(\n", " provider=\"palm\",\n", @@ -161,17 +127,9 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " Gravity is a force that pulls objects towards each other. It is what keeps us on the ground and keeps the planets in orbit around the sun. Gravity is always pulling on us, but we don't notice it because we are used to it. But if you jump up in the air, you will feel the force of gravity pulling you back down to the ground.\n" - ] - } - ], + "outputs": [], "source": [ "palm_gen = llm(\n", " provider=\"palm\",\n", @@ -196,17 +154,9 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Gravity is a force that pulls things together. It is what makes things fall to the ground and what holds us on the earth. Gravity is a fundamental force of nature that affects everything around us. It is a property of all matter, and it is what makes things heavy. Gravity is also what causes the moon to orbit the earth and the planets to orbit the sun. It is a very important force that plays a big role in our lives.\n" - ] - } - ], + "outputs": [], "source": [ "cohere_gen = llm(provider=\"cohere\")\n", "cohere_config = TextGenerationConfig(model=\"command\", max_tokens=4050, use_cache=True)\n", @@ -223,43 +173,110 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/victordibia/miniconda3/envs/llmx/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n", - "/home/victordibia/miniconda3/envs/llmx/lib/python3.11/site-packages/transformers/utils/hub.py:124: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.\n", - " warnings.warn(\n", - "Loading checkpoint shards: 100%|██████████| 8/8 [00:06<00:00, 1.27it/s]\n" - ] - } - ], + "outputs": [], "source": [ "hf_generator = llm(provider=\"hf\", model=\"HuggingFaceH4/zephyr-7b-beta\", device_map=\"auto\")" ] }, { "cell_type": "code", - "execution_count": 11, + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "hf_config = TextGenerationConfig(temperature=0, max_tokens=650, use_cache=False)\n", + "hf_response = hf_generator.generate(messages, config=hf_config)\n", + "print(hf_response.text[0].content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Ollama" + ] + }, + { + "cell_type": "code", + "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "Gravity is a special kind of force that pulls things down towards the ground. It's what makes apples fall from trees and why we don't float away into space! Gravity is also what keeps the Earth spinning around and around, so we don't fall off! It's a very strong force that we can't see, but we can feel it pulling us down when we jump into the air. Gravity is a very important force that helps keep everything in the universe in its place!\n" + "\n", + "Okay, so I need to figure out what gravity is. From what I remember in school, it's a force that pulls things towards the ground or something like that. But let me try to think through it more carefully.\n", + "\n", + "First, when we jump up and then come back down, that must be because of gravity. Without gravity, we'd just float away. So gravity is definitely related to keeping us grounded here on Earth. But I'm not exactly sure how it works. Is it a force? I think so, but what kind?\n", + "\n", + "I recall learning about Isaac Newton and his law of gravitation. Maybe that's where it starts. Newton said something about every object in the universe attracting every other object with a force proportional to their masses and inversely proportional to the distance between them squared. So F = G*(m1*m2)/r², right? Where F is the force, m1 and m2 are the masses, r is the distance, and G is the gravitational constant.\n", + "\n", + "But wait, isn't there also Einstein's theory of general relativity that explains gravity differently? I think he described it as space-time curvature caused by mass and energy. So, instead of being a force in the traditional sense, maybe gravity is a result of how massive objects bend the fabric of spacetime around them, and other objects move along these curves.\n", + "\n", + "So which one is correct? Is gravity both a force from Newton's perspective and also described as space-time curvature by Einstein? It seems like they complement each other. Newtonian gravity works well for everyday situations on Earth, but Einstein's theory is more accurate for stronger gravitational fields or when dealing with very massive objects like black holes.\n", + "\n", + "Let me think about how this affects things. For example, the orbit of planets around the Sun is explained by both theories, I believe. Newton could explain Kepler's laws using his law of gravitation, and Einstein provided a more precise description, especially for Mercury's orbit, which Newtonian gravity couldn't fully account for.\n", + "\n", + "But then there are practical applications. Satellites in orbit use Newtonian calculations because the effects of general relativity are small enough for their purposes, right? So maybe in most engineering and physics problems we don't need to consider space-time curvature unless we're dealing with extremely precise measurements or very massive objects like neutron stars or black holes.\n", + "\n", + "So putting it all together: Gravity is a natural phenomenon that causes all objects with mass or energy to exert force on each other. It can be described by Newton's law as a force proportional to the product of masses and inversely proportional to the square of the distance between their centers. However, Einstein's theory offers a more comprehensive explanation involving the curvature of space-time caused by massive objects.\n", + "\n", + "I wonder if gravity has any other effects besides attracting objects. Like, does it influence time? I think gravitational time dilation is a thing where time passes slower in stronger gravitational fields. So clocks near a massive object tick slower compared to those far away. That's an effect predicted by general relativity and hasn't been directly tested for everyday objects, but it's significant in certain contexts like GPS satellites which have to account for both special and general relativistic effects.\n", + "\n", + "Also, black holes are extreme cases of gravity where the escape velocity exceeds the speed of light, so not even light can escape. The event horizon is the boundary beyond which nothing can escape, including information about what's inside.\n", + "\n", + "But I'm still a bit confused about how exactly mass affects this force. If two objects with different masses interact via gravity, like a planet and a satellite, does each exert an equal force on the other? Newton's third law says yes, action equals reaction. So both the planet and satellite pull on each other with forces of equal magnitude but opposite direction.\n", + "\n", + "Another thought: what about electromagnetic forces compared to gravity? Gravity is by far the weakest force we experience daily, except for in extreme situations like black holes or neutron stars where densities are so high that gravitational effects dominate over electromagnetic ones. That's why particles can be held together despite their mutual repulsion due to charges; it's the strong gravitational pull keeping them from flying apart.\n", + "\n", + "Wait, but if gravity is a force, how do objects move? For example, when you throw a ball up, it decelerates because Earth's gravity pulls it down. So acceleration towards the center is what defines an object in free fall. But in space, without air resistance and other forces, objects would just keep moving at constant velocity or accelerate due to gravity.\n", + "\n", + "I'm also thinking about how gravity affects light. Since light has energy, Einstein said it should bend around massive objects because of curved spacetime. This phenomenon is gravitational lensing, where the path of light from a distant object is bent by a massive object in front, like a galaxy cluster. That's why we can see very distant galaxies sometimes.\n", + "\n", + "But wait, isn't that also related to gravity waves? Those are disturbances in spacetime caused by massive accelerating objects, like merging black holes or neutron stars. Gravitational waves carry energy away from systems, causing them to lose rotational energy, which is why some binary pulsars show their orbital period decreasing over time due to gravitational wave emission.\n", + "\n", + "So summarizing my thoughts: Gravity is a fundamental force of nature that acts between any two objects with mass. It can be described classically by Newton's law as a mutual attractive force proportional to the product of masses and inversely proportional to the square of the distance. Einstein's general relativity provides a more precise description involving the curvature of spacetime caused by mass-energy, which explains phenomena like gravitational time dilation, lensing, and the existence of black holes.\n", + "\n", + "It's also important to note that while gravity is weak on a human scale compared to other forces, it becomes significant at large scales or with very massive objects. This force is essential for structure formation in the universe, holding galaxies together, planets orbiting stars, and so on.\n", + "\n", + "I might have missed some aspects or mixed up some details, but this seems like a reasonable understanding based on what I've learned so far.\n", + "\n", + "\n", + "Gravity is one of the fundamental forces of nature that describes the mutual attraction between all objects with mass or energy. It can be understood through two primary frameworks:\n", + "\n", + "1. **Newtonian Gravity**: Formulated by Sir Isaac Newton, it states that every particle attracts every other particle with a force proportional to the product of their masses and inversely proportional to the square of the distance between their centers. This is expressed mathematically as \\( F = G \\frac{m_1 m_2}{r^2} \\), where \\( F \\) is the gravitational force, \\( m_1 \\) and \\( m_2 \\) are the masses of the objects, \\( r \\) is the distance between their centers, and \\( G \\) is the gravitational constant.\n", + "\n", + "2. **General Relativity**: Albert Einstein's theory describes gravity not as a force, but as the curvature of spacetime caused by mass and energy. In this framework, objects move along geodesics (the shortest paths) in curved spacetime, which gives the appearance of being \"attracted\" to massive objects.\n", + "\n", + "Gravity has several key effects and phenomena:\n", + "- **Planetary Motion**: It explains Kepler's laws and the orbits of planets around stars.\n", + "- **Gravitational Time Dilation**: Clocks near massive objects tick slower due to the curvature of spacetime, a concept tested in GPS satellite operations.\n", + "- **Black Holes**: Extreme gravitational forces where light cannot escape, characterized by event horizons.\n", + "- **Gravitational Lensing**: Light from distant objects bends around massive objects like galaxy clusters.\n", + "- **Gravitational Waves**: Disturbances in spacetime caused by accelerating masses, such as merging black holes.\n", + "\n", + "Gravity is crucial for structure formation in the universe, maintaining the cohesion of celestial bodies and their orbits. While it is the weakest force on a human scale, its effects are significant at large scales or with massive objects like black holes and neutron stars.\n", + "\n", + "In summary, gravity is a fundamental force described classically by Newton's law and more precisely by Einstein's general relativity, influencing everything from planetary motion to the behavior of massive cosmic structures.\n" ] } ], "source": [ - "hf_config = TextGenerationConfig(temperature=0, max_tokens=650, use_cache=False)\n", - "hf_response = hf_generator.generate(messages, config=hf_config)\n", - "print(hf_response.text[0].content)" + "ollama_generator = llm(provider=\"ollama\")\n", + "ollama_config = TextGenerationConfig(model=\"deepseek-r1:latest\", use_cache=False)\n", + "ollama_response = ollama_generator.generate(messages, config=ollama_config)\n", + "print(ollama_response.text[0].content)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": {