model load class updated

app.py

@@ -1,24 +1,17 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoPeftModelForCausalLM
 import gradio as gr
 
-# Load the base model and tokenizer
-base_model_path = "NousResearch/Llama-2-7b-chat-hf"  # Path to the base model
-tokenizer_path = "BoburAmirov/test-llama-uz"  # Path to the tokenizer
+# Load the fine-tuned model and tokenizer
+model_path = "BoburAmirov/test-llama-uz"  # Adjust this to the path where your fine-tuned model is saved
 
-# Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+model = AutoPeftModelForCausalLM.from_pretrained(model_path, device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+# Ensure the tokenizer settings match those used during training
 tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = "right"
 
-# Load the base model
-base_model = AutoModelForCausalLM.from_pretrained(base_model_path)
-
-# Load the adapter
-adapter_path = "BoburAmirov/test-llama-uz/adapter_model.safetensors"
-model = PeftModel.from_pretrained(base_model, adapter_path)
-
 # Set the model to evaluation mode
 model.eval()
 
@@ -30,7 +23,7 @@ def generate_text(input_prompt):
     with torch.no_grad():
         output = model.generate(
             input_ids,
-            max_length=200,  # Adjust max_length as needed
+            max_length=400,  # Adjust max_length as needed
             num_return_sequences=1,
             temperature=0.7,  # Control randomness
             top_p=0.9,  # Control diversity
@@ -42,16 +35,13 @@ def generate_text(input_prompt):
     return generated_text
 
 # Create a Gradio interface
-interface = gr.Interface(
+iface = gr.Interface(
     fn=generate_text,
     inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
     outputs="text",
-    title="Text Generation with LLaMA-2",
-    description="Enter a prompt and get generated text from the fine-tuned LLaMA-2 model."
+    title="Text Generation with LLaMA",
+    description="Generate text using a fine-tuned LLaMA model."
 )
 
-# Launch the Gradio interface
-
-
 if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860)