fix linting errors

HunterGerlach · Oct 5, 2023 · 9237e73 · 9237e73
1 parent a3cea80
commit 9237e73
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/src/hosted_llm.py b/src/hosted_llm.py
@@ -1,5 +1,7 @@
-import requests
+"""Module for handling self-hosted LLama2 models"""
+
 from typing import Any, List, Mapping, Optional
+import requests
 from langchain.callbacks.manager import CallbackManagerForLLMRun
 from langchain.llms.base import LLM
 from langchain.schema.output_parser import BaseOutputParser
@@ -35,16 +37,15 @@ def _identifying_params(self) -> Mapping[str, Any]:
         """Get the identifying parameters."""
         return {"uri": self.uri}
 
-
-class CustomLlamaParser(BaseOutputParser[str]):
+class CustomLlamaParser(BaseOutputParser[str]): # pylint: disable=R0903
     """Class to correctly parse model outputs"""
 
     def parse(self, text:str) -> str:
+        """Parse the output of our LLM"""
         if text.startswith("Model Server is not Working due"):
             return text
-        else:
-            cleaned = str(text).split("[/INST]")
-            return cleaned[1]
+        cleaned = str(text).split("[/INST]")
+        return cleaned[1]
 
     @property
     def _type(self) -> str:

diff --git a/src/v1/endpoints.py b/src/v1/endpoints.py
@@ -5,12 +5,12 @@
 from pydantic import BaseModel # pylint: disable=E0611
 from langchain.llms import OpenAI
 from langchain.llms import VertexAI
-from src.hosted_llm import HostedLLM
-from src.hosted_llm import CustomLlamaParser
 from langchain.callbacks import get_openai_callback
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 
+from src.hosted_llm import HostedLLM
+from src.hosted_llm import CustomLlamaParser
 from src.config import Config
 from src.embeddings import EmbeddingSource
 from src.logging_setup import setup_logger