Update app.py

app.py

@@ -1,24 +1,26 @@
 import gradio as gr
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from peft import PeftModel
 
-# Load the model and tokenizer
-model_name = "raccoote/angry-birds-v2"  # Replace with the correct model name
+# Load the base model and tokenizer directly from the raccoote repository
+model_name = "raccoote/angry-birds-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+base_model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Load the LoRA adapter from the raccoote repository
+adapter_model = PeftModel.from_pretrained(base_model, model_name)
 
 def generate_text(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    # Process the outputs to generate text (this will vary based on your model)
-    # Here we just return the hidden states shape as a placeholder
-    return outputs.last_hidden_state.shape
+    outputs = adapter_model.generate(**inputs, max_new_tokens=50)
+    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 Model with LoRA Adapters",
+                     title="LLaMA 3.1 with LoRA Adapters",
                      description="Enter a prompt and get the model's output.")
 
 iface.launch()