Update app.py

app.py

@@ -1,5 +1,14 @@
 import fitz  # PyMuPDF for PDF extraction
 import re
+import unsloth
+import os
+from huggingface_hub import login
+from datasets import Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+from peft import LoraConfig, get_peft_model
+import gradio as gr
+from transformers import pipeline
+
 
 def extract_text_from_pdf(pdf_path):
     """Extract text from a PDF file"""
@@ -14,8 +23,6 @@ def preprocess_text(text):
 pdf_text = extract_text_from_pdf("new-american-standard-bible.pdf")
 clean_text = preprocess_text(pdf_text)
 
-import os
-from huggingface_hub import login
 
 # Read the Hugging Face token from environment variables
 hf_token = os.getenv("access_token")
@@ -26,12 +33,8 @@ if hf_token is None:
 # Log in to Hugging Face
 login(token=hf_token)
 
-from datasets import Dataset
-from transformers import AutoTokenizer
-import unsloth
-
-
-model_name = "meta-llama/Llama-2-7b-hf"  # You can use a smaller one like "meta-llama/Llama-2-7b-chat-hf"
+#model_name = "meta-llama/Llama-2-7b-hf"  # You can use a smaller one like "meta-llama/Llama-2-7b-chat-hf"
+model_name = "unsloth/llama-2-7b-chat"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -45,7 +48,6 @@ tokenizer.pad_token = tokenizer.eos_token  # Use EOS as PAD token
 # tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
 # Tokenization function
-@unsloth.optimize
 def tokenize_function(examples):
     tokens = tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
     tokens["labels"] = tokens["input_ids"].copy()  # Use input as labels for text generation
@@ -53,9 +55,6 @@ def tokenize_function(examples):
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-from transformers import AutoModelForCausalLM, TrainingArguments, Trainer
-from peft import LoraConfig, get_peft_model
-
 # Load LLaMA 2 model in 4-bit mode to save memory
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -103,9 +102,8 @@ perform_training()
 model.save_pretrained("./fine_tuned_llama2")
 tokenizer.save_pretrained("./fine_tuned_llama2")
 
-import gradio as gr
-from transformers import pipeline
 
+# CHATBOT START
 chatbot = pipeline("text-generation", model="./fine_tuned_llama2")
 
 def chatbot_response(prompt):