Update app.py

app.py

@@ -26,34 +26,40 @@ vectordb = Chroma(
 
 retriever = vectordb.as_retriever(search_kwargs={"k": 5})
 
-from typing import Optional, List, Dict, Any
-# LLM, KwArgsModel, InferenceClient 등 필요한 모듈 또는 클래스를 불러와야 함
+class KwArgsModel(BaseModel):
+    kwargs: Dict[str, Any] = Field(default_factory=dict)
 
-class CustomInferenceClient:
-    def __init__(self, model_name: str, hf_token: str, kwargs: Optional[Dict[str, Any]] = None):
-        self.llm = LLM(model_name=model_name, hf_token=hf_token)  # LLM 인스턴스 생성
-        self.kw_args_model = KwArgsModel(kwargs=kwargs)  # KwArgsModel 인스턴스 생성
-        self.inference_client = InferenceClient(model=model_name, token=hf_token)  # InferenceClient 인스턴스 생성
-        self.model_name = model_name  # 모델 이름 저장
+class CustomInferenceClient(LLM, KwArgsModel):
+    model_name: str
+    inference_client: InferenceClient
 
-    # _call 메서드 구현
-    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+    def __init__(self, model_name: str, hf_token: str, kwargs: Optional[Dict[str, Any]] = None):
+        inference_client = InferenceClient(model=model_name, token=hf_token)
+        super().__init__(
+            model_name=model_name,
+            hf_token=hf_token,
+            kwargs=kwargs,
+            inference_client=inference_client
+        )
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None
+    ) -> str:
         if stop is not None:
-            raise ValueError("stop kwargs are not permitted.")  # stop 인자가 있으면 에러 발생
-        kwargs = self.kw_args_model.kwargs  # KwArgsModel로부터 kwargs 받기
-        response_gen = self.inference_client.text_generation(prompt, **kwargs, stream=True)  # 텍스트 생성 요청
-        response = ''.join(response_gen)  # 스트림으로부터 문자열 생성
-        return response  # 생성된 문자열 반환
+            raise ValueError("stop kwargs are not permitted.")
+        response_gen = self.inference_client.text_generation(prompt, **self.kwargs, stream=True)
+        response = ''.join(response_gen)  
+        return response
 
-    # _llm_type 속성 구현
     @property
     def _llm_type(self) -> str:
-        return "custom"  # 사용자 정의 타입으로 지정
+        return "custom"
 
-    # _identifying_params 속성 구현
     @property
     def _identifying_params(self) -> dict:
-        return {"model_name": self.model_name}  # 모델 이름을 반환
+        return {"model_name": self.model_name}
 
 kwargs = {"max_new_tokens":256, "temperature":0.9, "top_p":0.6, "repetition_penalty":1.3, "do_sample":True}