Update app.py

app.py

@@ -13,7 +13,7 @@ from datetime import datetime
 import json
 import gradio as gr
 import re 
-
+from unsloth import FastLanguageModel
 class DocumentRetrievalAndGeneration:
     def __init__(self, embedding_model_name, lm_model_id, data_folder):
         # hf_token = os.getenv('HF_TOKEN')
@@ -66,15 +66,9 @@ class DocumentRetrievalAndGeneration:
                 
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         model = AutoModelForCausalLM.from_pretrained(model_id)
-        generate_text = pipeline(
-            model=model,
-            tokenizer=tokenizer,
-            return_full_text=True,
-            task='text-generation',
-            temperature=0.6,
-            max_new_tokens=256,
-        )
-        return generate_text
+        FastLanguageModel.for_inference(model) # Enable native 2x faster inference
+        
+        # return generate_text
 
     def generate_response_with_timeout(self, model_inputs):
         try:
@@ -127,16 +121,28 @@ class DocumentRetrievalAndGeneration:
         </s>
         """
 
-        messages = [{"role": "user", "content": prompt}]
-        encodeds = self.llm.tokenizer.apply_chat_template(messages, return_tensors="pt")
-        model_inputs = encodeds.to(self.llm.device)
+        # messages = [{"role": "user", "content": prompt}]
+        # encodeds = self.llm.tokenizer.apply_chat_template(messages, return_tensors="pt")
+        # model_inputs = encodeds.to(self.llm.device)
         
-        start_time = datetime.now()
-        generated_ids = self.generate_response_with_timeout(model_inputs)
-        elapsed_time = datetime.now() - start_time
+        # start_time = datetime.now()
+        # generated_ids = self.generate_response_with_timeout(model_inputs)
+        # elapsed_time = datetime.now() - start_time
+
+        # decoded = self.llm.tokenizer.batch_decode(generated_ids)
+        # generated_response = decoded[0]
 
-        decoded = self.llm.tokenizer.batch_decode(generated_ids)
-        generated_response = decoded[0]
+        inputs = tokenizer(
+        [
+            alpaca_prompt.format(
+                "", # instruction
+                prompt, # input
+                "", # output - leave this blank for generation!
+            )
+        ], return_tensors = "pt")#.to("cuda")
+        
+        outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True)
+        tokenizer.batch_decode(outputs)
         match1 = re.search(r'\[/INST\](.*?)</s>', generated_response, re.DOTALL)
         
         match2 = re.search(r'Solution:(.*?)</s>', generated_response, re.DOTALL | re.IGNORECASE)