Update app.py

app.py

@@ -1,31 +1,49 @@
-import gradio as gr
-import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel
+import gradio as gr
+
+# Step 1: Load the base model
+base_model_name = "meta-llama/Llama-3.3-70B-Instruct"
+adapter_repo = "daresearch/Llama-3.3-70B-ft-exec-roles"
 
-BASE_MODEL = "meta-llama/Llama-3.3-70B-Instruct"
-ADAPTER = "daresearch/Llama-3.3-70B-ft-exec-roles"
+# Load the base model (LlamaForCausalLM)
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    device_map="auto",           # Distribute model across GPUs (if available)
+    torch_dtype=torch.float16,   # Use FP16 precision to save memory
+)
+
+# Load the LoRA adapter into the base model
+model_with_adapter = PeftModel.from_pretrained(
+    base_model,
+    adapter_repo,
+    device_map="auto",
+)
 
-# Load the base model
-base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.float16, device_map="auto")
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+# Extract the underlying base model for compatibility with pipelines
+underlying_model = model_with_adapter.base_model
 
-# Load the adapter
-model = PeftModel.from_pretrained(base_model, ADAPTER, device_map="auto")
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 
-# Create a text generation pipeline
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# Create the text generation pipeline
+pipe = pipeline("text-generation", model=underlying_model, tokenizer=tokenizer)
 
-def infer(prompt):
-    outputs = pipe(prompt, max_length=128)
+# Define the Gradio interface function
+def generate_text(prompt):
+    # Use the pipeline to generate text
+    outputs = pipe(prompt, max_length=200)
     return outputs[0]["generated_text"]
 
+# Create the Gradio interface
 iface = gr.Interface(
-    fn=infer,
+    fn=generate_text,
     inputs="text",
     outputs="text",
-    title="LoRA-Enhanced Model"
+    title="LoRA-Enhanced LLaMA Text Generator",
+    description="Provide a prompt, and the model will generate a response."
 )
 
+# Launch the app
 if __name__ == "__main__":
     iface.launch()