Update app.py

app.py

@@ -26,6 +26,8 @@ from transformers import (
     Trainer
 )
 
+from peft import PeftModel
+
 from datasets import Dataset
 from accelerate import Accelerator
 # Imports specific to the custom peft lora model
@@ -105,10 +107,39 @@ class_weights = compute_class_weight(class_weight='balanced', classes=classes, y
 accelerator = Accelerator()
 class_weights = torch.tensor(class_weights, dtype=torch.float32).to(accelerator.device)
 
-dubug_result = class_weights
-demo = gr.Blocks(title="DEMO FOR ESMBind")
+# inference
+# Path to the saved LoRA model
+model_path = "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_v2_cp3"
+# ESM2 base model
+base_model_path = "facebook/esm2_t12_35M_UR50D"
+
+# Load the model
+base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
+loaded_model = PeftModel.from_pretrained(base_model, model_path)
+
+# Ensure the model is in evaluation mode
+loaded_model.eval()
+
+# Protein sequence for inference
+protein_sequence = "MAVPETRPNHTIYINNLNEKIKKDELKKSLHAIFSRFGQILDILVSRSLKMRGQAFVIFKEVSSATNALRSMQGFPFYDKPMRIQYAKTDSDIIAKMKGT"  # Replace with your actual sequence
+
+# Tokenize the sequence
+inputs = tokenizer(protein_sequence, return_tensors="pt", truncation=True, max_length=1024, padding='max_length')
+
+# Run the model
+with torch.no_grad():
+    logits = loaded_model(**inputs).logits
+
+# Get predictions
+tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])  # Convert input ids back to tokens
+predictions = torch.argmax(logits, dim=2)
+
+# debug result
+dubug_result = predictions  #class_weights
+
+demo = gr.Blocks(title="DEMO FOR ESM2Bind")
 
 with demo:
-    gr.Markdown("# DEMO FOR ESMBind")
+    gr.Markdown("# DEMO FOR ESM2Bind")
     gr.Textbox(dubug_result)
 demo.launch()