Update app.py

app.py

@@ -43,30 +43,30 @@ def create_llm():
     """Initialize the language model with optimized parameters"""
     
     
-    # bnb_config = BitsAndBytesConfig(
-    #     load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
-    # )
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
+    )
     
-    # model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, 
-    #                                              # quantization_config=bnb_config
-    #                                             )
-    # tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, 
+                                                 quantization_config=bnb_config
+                                                )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     
-    # terminators = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]
+    terminators = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]
     
-    # text_generation_pipeline = pipeline(
-    #     model=model,
-    #     tokenizer=tokenizer,
-    #     task="text-generation",
-    #     temperature=0.2,
-    #     do_sample=True,
-    #     repetition_penalty=1.1,
-    #     return_full_text=False,
-    #     max_new_tokens=200,
-    #     eos_token_id=terminators,
-    # )
+    text_generation_pipeline = pipeline(
+        model=model,
+        tokenizer=tokenizer,
+        task="text-generation",
+        temperature=0.2,
+        do_sample=True,
+        repetition_penalty=1.1,
+        return_full_text=False,
+        max_new_tokens=200,
+        eos_token_id=terminators,
+    )
     
-    # llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
+    llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
 
     # return HuggingFaceHub(
     #     repo_id=MODEL_NAME,
@@ -77,13 +77,13 @@ def create_llm():
     #         "frequency_penalty": 0.5
     #     }
     # )
-    llm = HuggingFaceEndpoint(
-            repo_id=MODEL_NAME,
-            huggingfacehub_api_token = api_token,
-            temperature = 0.7,
-            max_new_tokens = 1024,
-            top_k = 0.9,
-        )
+    # llm = HuggingFaceEndpoint(
+    #         repo_id=MODEL_NAME,
+    #         huggingfacehub_api_token = api_token,
+    #         temperature = 0.7,
+    #         max_new_tokens = 1024,
+    #         top_k = 0.9,
+    #     )
     # llm = HuggingFacePipeline.from_model_id(
     #     model_id=MODEL_NAME,
     #     task="text-generation",