Update README.md

README.md

@@ -19,6 +19,10 @@ An Evolutionary-scale Model (ESM) for protein function calling from amino acid s
 
 **Note**: This model specilizes on the `celluar component` subgraph of the gene ontology.
 
+## Code Repository
+
+https://github.com/andrewdalpino/esm2-function-classifier
+
 ## Model Specs
 
 **Embedding Dimensions**: 640
@@ -44,11 +48,13 @@ model = EsmForSequenceClassification.from_pretrained(model_name)
 
 model.eval()
 
-sequnce = "MCNAWYISVDFEKNREDKSKCIHTRRNSGPKLLEHVMYEVLRDWYCLEGENVYMMGKKWQMPMCSLH"
+sequence = "MCNAWYISVDFEKNREDKSKCIHTRRNSGPKLLEHVMYEVLRDWYCLEGENVYMMGKKWQMPMCSLH"
+
+top_k = 10
 
 out = tokenizer(
     sequence,
-    max_length=args.context_length,
+    max_length=1026,
     truncation=True,
 )
 
@@ -61,7 +67,7 @@ with torch.no_grad():
 
     probabilities = torch.sigmoid(outputs.logits.squeeze(0))
 
-    probabilities, indices = torch.topk(probabilities, args.top_k)
+    probabilities, indices = torch.topk(probabilities, top_k)
 
 probabilities = probabilities.tolist()