Update model.py

model.py

@@ -1,3 +1,4 @@
+
 import os
 import subprocess
 from dotenv import load_dotenv
@@ -8,16 +9,36 @@ try:
 except:
     PINECONE_API_KEY = subprocess.check_output(["bash", "-c", "echo ${{ secrets.PINECONE_API_KEY }}"]).decode("utf-8").strip()
 
-
+from typing import Optional,List,Mapping,Any
 from langchain.embeddings import HuggingFaceEmbeddings
 import pinecone
 import torch
 from langchain import PromptTemplate, LLMChain,HuggingFacePipeline
 from langchain.vectorstores import Pinecone
-from langchain.chains.question_answering import load_qa_chain
-from langchain.chains import RetrievalQA
+from langchain.llms.base import LLM
 from transformers import pipeline
 
+class CustomLLM(LLM):
+    # def __init__(self,model_name,pipeline):
+    model_name ="databricks/dolly-v2-3b"
+    num_output = 128
+    pipeline = pipeline(model=model_name, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto", 
+        return_full_text=True, do_sample=False, max_new_tokens=128)
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        prompt_length = len(prompt)
+        response = self.pipeline(prompt, max_new_tokens=self.num_output)[0]["generated_text"]
+
+        # only return newly generated tokens
+        return response[prompt_length:]
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        return {"name_of_model": self.model_name}
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom"
+    
 def get_llm(model_name,pinecone_index,llm):
     # model_name = "bert-large-uncased" #"t5-large"
     model_kwargs = {'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
@@ -31,7 +52,7 @@ def get_llm(model_name,pinecone_index,llm):
     )
 
     index = pinecone.Index(pinecone_index)
-    print(index.describe_index_stats())
+    # print(index.describe_index_stats())
 
     docsearch = Pinecone(index, embeddings.embed_query,"text")