Create custom_llm.py

custom_llm.py

@@ -0,0 +1,71 @@
+from typing import Any, List, Mapping, Optional
+from langchain_core.language_models.llms import LLM
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+
+from typing import Literal
+
+import requests
+
+
+class CustomLLM(LLM):
+    repo_id : str
+    api_token : str
+    model_type: Literal["text2text-generation", "text-generation"]
+    max_new_tokens: int = None
+    temperature: float = 0.001
+    timeout: float = None
+    top_p: float = None
+    top_k : int = None
+    repetition_penalty : float = None
+    stop : List[str] = []
+
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+
+        headers = {"Authorization": f"Bearer {self.api_token}"}
+        API_URL = f"https://api-inference.huggingface.co/models/{self.repo_id}"
+
+        parameters_dict = {
+          'max_new_tokens': self.max_new_tokens,
+          'temperature': self.temperature,
+          'timeout': self.timeout,
+          'top_p': self.top_p,
+          'top_k': self.top_k,
+          'repetition_penalty': self.repetition_penalty,
+          'stop':self.stop
+        }
+
+        if self.model_type == 'text-generation':
+          parameters_dict["return_full_text"]=False
+
+        data = {"inputs": prompt, "parameters":parameters_dict, "options":{"wait_for_model":True}}
+        data = requests.post(API_URL, headers=headers, json=data).json()
+        try:
+            return data[0]['generated_text']
+        except:
+            return data
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        return  {
+          'repo_id': self.repo_id,
+          'model_type':self.model_type,
+          'stop_sequences':self.stop,
+          'max_new_tokens': self.max_new_tokens,
+          'temperature': self.temperature,
+          'timeout': self.timeout,
+          'top_p': self.top_p,
+          'top_k': self.top_k,
+          'repetition_penalty': self.repetition_penalty
+      }