Spaces:
Running
Running
| """The code borrowed from https://colab.research.google.com/drive/1RW2yTxh5b9w7F3IrK00Iz51FTO5W01Rx?usp=sharing#scrollTo=RgbLVmf-o4j7""" | |
| import os | |
| from typing import Any, Dict | |
| import together | |
| from pydantic import Extra, root_validator | |
| from langchain.llms.base import LLM | |
| from langchain.utils import get_from_dict_or_env | |
| from toolkit.utils import Config | |
| configs = Config("configparser.ini") | |
| os.environ["TOGETHER_API_KEY"] = configs.together_api_key | |
| # together.api_key = configs.together_api_key | |
| # models = together.Models.list() | |
| # for idx, model in enumerate(models): | |
| # print(idx, model["name"]) | |
| class TogetherLLM(LLM): | |
| """Together large language models.""" | |
| model: str = "togethercomputer/llama-2-70b-chat" | |
| """model endpoint to use""" | |
| together_api_key: str = os.environ["TOGETHER_API_KEY"] | |
| """Together API key""" | |
| temperature: float = 0 | |
| """What sampling temperature to use.""" | |
| max_tokens: int = 512 | |
| """The maximum number of tokens to generate in the completion.""" | |
| class Config: | |
| extra = "forbid" | |
| # @root_validator() | |
| # def validate_environment(cls, values: Dict) -> Dict: | |
| # """Validate that the API key is set.""" | |
| # api_key = get_from_dict_or_env(values, "together_api_key", "TOGETHER_API_KEY") | |
| # values["together_api_key"] = api_key | |
| # return values | |
| def _llm_type(self) -> str: | |
| """Return type of LLM.""" | |
| return "together" | |
| def _call( | |
| self, | |
| prompt: str, | |
| **kwargs: Any, | |
| ) -> str: | |
| """Call to Together endpoint.""" | |
| together.api_key = self.together_api_key | |
| output = together.Complete.create( | |
| prompt, | |
| model=self.model, | |
| max_tokens=self.max_tokens, | |
| temperature=self.temperature, | |
| ) | |
| text = output["output"]["choices"][0]["text"] | |
| return text | |
| # if __name__ == "__main__": | |
| # test_llm = TogetherLLM( | |
| # model="togethercomputer/llama-2-70b-chat", temperature=0, max_tokens=1000 | |
| # ) | |
| # print(test_llm("What are the olympics? ")) | |