Update app.py

app.py

@@ -1,35 +1,35 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
+import torch
 
-# Define the base model and configuration
-base_model_name = "raccoote/angry-birds-v2"
+# Load the base model and tokenizer
+model_id = "unsloth/Meta-Llama-3.1-8B"  # Use the appropriate LLaMA 3.1 8b model ID
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)  # Use torch.float32 for CPU
+model.to("cpu")  # Ensure the model is loaded on CPU
 
-# Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+# Load your LoRA adapter
+adapter_repo = "raccoote/angry-birds-v2"  # Your repository path
+adapter_weight_name = "adapter_model.safetensors"  # The weight file name
 
-# Load the model with 8-bit precision
-quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+# Load LoRA weights
+peft_model = PeftModel.from_pretrained(model, adapter_repo, weight_name=adapter_weight_name, adapter_name="angry_birds")
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_name,
-    quantization_config=quantization_config,
-    device_map="auto"  # This will ensure the model is distributed to available hardware
-)
-
-# Load the LoRA adapter from the repository
-adapter_model = PeftModel.from_pretrained(base_model, base_model_name)
-
-def generate_text(prompt):
+# Prepare for inference
+def generate_text(prompt, model, tokenizer, peft_model, max_length=50):
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = adapter_model.generate(**inputs, max_new_tokens=50)
+    outputs = peft_model.generate(
+        **inputs, 
+        max_length=max_length, 
+        num_return_sequences=1,
+        do_sample=True,  # or use `do_sample=False` for deterministic outputs
+        top_p=0.95,  # or other sampling parameters
+        temperature=0.7
+    )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Create the Gradio interface
-iface = gr.Interface(fn=generate_text, 
-                     inputs="text", 
-                     outputs="text",
-                     title="LLaMA 3.1 with LoRA Adapters",
-                     description="Enter a prompt and get the model's output.")
+# Generate text with the loaded LoRA adapter
+prompt = "large piggy on wooden tower"
+generated_text = generate_text(prompt, model, tokenizer, peft_model)
 
-iface.launch()
+print(generated_text)