Create handler.py

handler.py

@@ -0,0 +1,58 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.llms import HuggingFacePipeline
+from langchain.chains import RetrievalQA
+import torch
+
+class Handler:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {self.device}")
+
+        # Load the fine-tuned model and tokenizer
+        print("Loading model and tokenizer...")
+        self.model = AutoModelForCausalLM.from_pretrained("PranavKeshav/upf_code_generator").to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained("PranavKeshav/upf_code_generator").to(self.device)
+
+        # Load the FAISS index and embeddings
+        print("Loading FAISS index and embeddings...")
+        self.embeddings = HuggingFaceEmbeddings()
+        self.vectorstore = FAISS.load_local("faiss_index", self.embeddings, allow_dangerous_deserialization=True)
+
+        # Create the Hugging Face pipeline for text generation
+        print("Creating Hugging Face pipeline...")
+
+        def run_inference(prompt: str):
+            # Assuming 2048 is the desired max sequence length
+            return self.model.generate(
+                prompt, temperature=0.7, max_length=2048, top_p=0.95, repetition_penalty=1.15
+            )
+
+        self.hf_pipeline = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            temperature=0.7,
+            max_new_tokens=2048,
+            top_p=0.95,
+            repetition_penalty=1.15
+        )
+
+        self.hf_pipeline.model.generate = run_inference
+        # Wrap the pipeline in LangChain
+        self.llm = HuggingFacePipeline(pipeline=self.hf_pipeline)
+
+        # Create the retriever and pipeline
+        self.retriever = self.vectorstore.as_retriever()
+        self.qa_chain = RetrievalQA.from_chain_type(llm=self.llm, retriever=self.retriever)
+
+    def __call__(self, request):
+        # Get the prompt from the request
+        prompt = request.json.get("prompt")
+
+        # Generate UPF code using the QA chain
+        response = self.qa_chain.run(prompt)
+
+        # Return the response
+        return {"response": response}