Update app.py

app.py

@@ -62,6 +62,105 @@ prompt = ChatPromptTemplate.from_messages(
         ("human", human),
     ]
 )
+from typing import Any, Dict, List, Optional
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
+from langchain_core.outputs import ChatResult, ChatGeneration
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.callbacks.manager import AsyncCallbackManagerForLLMRun
+from langchain_core.runnables import run_in_executor
+from transformers import AutoProcessor, AutoModelForCausalLM
+import torch
+class Chatchat(BaseChatModel):
+
+    model_name: str = "peterpeter8585/deepseek_1"
+    tokenizer : AutoTokenizer = None
+    model: AutoModelForCausalLM = None
+    model_path: str = None
+
+    def __init__(self, model_path, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        if model_path is not None:
+            self.model_name = model_path
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name, trust_remote_code=True)
+
+
+    def _call(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> str:
+        # Load and preprocess the image
+        messages = [
+            {"role": "system", "content": "You are Chatchat.A helpful assistant at code."},
+            {"role": "user", "content": prompt}
+        ]
+
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        model_inputs  = self.tokenizer([text], return_tensors="pt").to(self.model.device)
+        generated_ids = self.model.generate(
+              **model_inputs,
+              max_new_tokens=512
+          )
+        generated_ids = [
+            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+          ]
+
+        response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        return response
+
+    async def _acall(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> str:
+        # Implement the async logic to generate a response from the model
+        return await run_in_executor(
+            None,
+            self._call,
+            prompt,
+            stop,
+            run_manager.get_sync() if run_manager else None,
+            **kwargs,
+        )
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom-llm-chat"
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        return {"model_name": self.model_name}
+
+    def _generate(
+            self,
+            messages: List[BaseMessage],
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> ChatResult:
+        # Assumes the first message contains the prompt and the image path is in metadata
+        prompt = messages[0].content
+        response_text = self._call(prompt,  stop, run_manager, **kwargs)
+
+        # Create AIMessage with the response
+        ai_message = AIMessage(content=response_text)
+        return ChatResult(generations=[ChatGeneration(message=ai_message)])
+
+
+llm=Chatchat()
 #from transformers import pipeline,AutoModelForCausalLM as M,AutoTokenizer as T
 #m=M.from_pretrained("peterpeter8585/syai4.3")
 #t=T.from_pretrained("peterpeter8585/syai4.3")