Update handler.py

handler.py

@@ -7,13 +7,10 @@ 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)
+        self.model = AutoModelForCausalLM.from_pretrained("anirudh248/upf_code_generator_final", device_map="auto")
+        self.tokenizer = AutoTokenizer.from_pretrained("anirudh248/upf_code_generator_final")
 
         # Load the FAISS index and embeddings
         print("Loading FAISS index and embeddings...")
@@ -22,37 +19,40 @@ class Handler:
 
         # 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,
+            device=0 if torch.cuda.is_available() else -1,
             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
+        # Create the retriever and RetrievalQA chain
         self.retriever = self.vectorstore.as_retriever()
-        self.qa_chain = RetrievalQA.from_chain_type(llm=self.llm, retriever=self.retriever)
+        self.qa_chain = RetrievalQA.from_chain_type(
+            llm=self.llm,
+            retriever=self.retriever,
+            return_source_documents=False
+        )
 
     def __call__(self, request):
-        # Get the prompt from the request
-        prompt = request.json.get("prompt")
+        try:
+            # Get the prompt from the request
+            prompt = request.json.get("prompt")
+            if not prompt:
+                return {"error": "Prompt is required"}, 400
+
+            # Generate UPF code using the QA chain
+            response = self.qa_chain.run(prompt)
 
-        # Generate UPF code using the QA chain
-        response = self.qa_chain.run(prompt)
+            # Return the response
+            return {"response": response}
 
-        # Return the response
-        return {"response": response}
+        except Exception as e:
+            return {"error": str(e)}, 500