Update app.py

app.py

@@ -4,21 +4,21 @@ import torch
 from transformers import AutoTokenizer, AutoModel
 from safetensors.torch import load_file
 
-# Load the Hugging Face API token
+# Load the Hugging Face API token from environment variable
 token = os.getenv("HUGGINGFACE_API_TOKEN")
 if not token:
     raise ValueError("HUGGINGFACE_API_TOKEN is not set. Please add it in the Secrets section of your Space.")
 
-# Configure device and data type
+# Configure device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the tokenizer and model
+# Load the tokenizer and model using the token
 model_repo = "Grandediw/lora_model"
-tokenizer = AutoTokenizer.from_pretrained(model_repo, use_auth_token=True)
-base_model = AutoModel.from_pretrained(model_repo, use_auth_token=True)
+tokenizer = AutoTokenizer.from_pretrained(model_repo, token=token)
+base_model = AutoModel.from_pretrained(model_repo, token=token)
 
-# Load LoRA adapter weights
-lora_weights_path = "adapter_model.safetensors"  # Ensure this file is present in the same directory
+# Load the LoRA adapter weights
+lora_weights_path = "adapter_model.safetensors"  # Ensure this file exists in the same directory
 lora_weights = load_file(lora_weights_path)
 
 # Apply LoRA weights to the base model
@@ -29,45 +29,24 @@ for name, param in base_model.named_parameters():
 # Move the model to the device
 base_model = base_model.to(device)
 
-# Inference function
-def infer(prompt, negative_prompt=None):
+# Define the inference function
+def infer(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     outputs = base_model(**inputs)
-    return outputs.last_hidden_state.mean(dim=1).cpu().detach().numpy()  # Placeholder return
+    # Placeholder return, modify based on your specific model task
+    return outputs.last_hidden_state.mean(dim=1).cpu().detach().numpy()
 
-# Gradio Interface
-css = """
-#interface-container {
-    margin: 0 auto;
-    max-width: 700px;
-    padding: 15px;
-    border-radius: 10px;
-    background-color: #f9f9f9;
-    box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.1);
-}
-#header {
-    text-align: center;
-    font-size: 1.5em;
-    font-weight: bold;
-    margin-bottom: 20px;
-    color: #333;
-}
-"""
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## LoRA Model Inference")
 
-with gr.Blocks(css=css) as demo:
-    with gr.Box(elem_id="interface-container"):
-        gr.Markdown("<div id='header'>LoRA Model Inference</div>")
-        
-        # Input for prompt and run button
+    with gr.Row():
         prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
-        run_button = gr.Button("Generate Output", variant="primary")
+        generate_button = gr.Button("Generate")
 
-        # Display output
-        output = gr.Textbox(label="Output")
+    output = gr.Textbox(label="Output")
 
-        # Connect button with inference
-        run_button.click(fn=infer, inputs=[prompt], outputs=[output])
+    generate_button.click(fn=infer, inputs=[prompt], outputs=[output])
 
-# Launch the app
 if __name__ == "__main__":
     demo.launch()