Views
No views yet
prompt & article text the model will produce geneId: go term descriptions).1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4# Define parameters
5max_input_size = 10240
6temperature = 0.9
7top_k = 50
8top_p = 0.95
9
10# Load model
11tokenizer = AutoTokenizer.from_pretrained("danjwarr/gene-go-annotator-1b-single-term")
12model = AutoModelForCausalLM.from_pretrained("danjwarr/gene-go-annotator-1b-single-term", torch_dtype=torch.float16)
13
14# Prepare prompt
15# Note: The model does not require this prompt at inference time
16# The article text alone is sufficient; this is kept to match training
17system_prompt = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>You are an assistant tasked with extracting the gene function and location from biological literature only about [species]. You should ignore anything which describes a different species. You should describe any Biological Process, Molecular Function or Cellular Component related to the Gene product that has been identified. These descriptions should be one or two full sentences, in the same detail as given in the Gene Ontology. Some examples of entries in the Gene Ontology are <examples><example>. \"Entry of a symbiont into the body, tissues, or cells of a host organism as part of the symbiont life cycle. The host is defined as the larger of the organisms involved in a symbiotic interaction.\" </example> <example>. \"The secretion of neuropeptides contained within a dense core vesicle by fusion of the granule with the presynaptic membrane, stimulated by a rise in cytosolic calcium ion concentration </example> <example>. A heterotetrameric protein complex that associates with replication origins, where it is required for the initiation of DNA replication, and with replication forks </example> <example> Combining with the neurotransmitter dopamine and activating adenylate cyclase via coupling to Gi/Go to initiate a change in cell activity.</example> <example> Any process that modulates the frequency, rate or extent of the growth of all or part of an organism so that it occurs at its proper speed, either globally or in a specific part of the organisms development. </example> </examples><|eot_id|>\n<|start_header_id|>user<|end_header_id|>\nIn the following article, look for evidence of any Biological Processes, molecular functions, or cellular components that involve [species] and the gene product identified by [geneID]. Give a quote from the article of the supporting evidence as complete sentences. Supply these as a numbered list. If there is no evidence of any specific biological process, molecular function, or cellular component, reply with an empty list []. Add no other commentary. The article follows: """
18
19# ***NEED INPUT HERE***
20article_text = "" # replace this with research article text
21
22# Create the full prompt
23prompt = system_prompt + article_text + " <TERMS>" # system_prompt may be omitted
24
25# Tokenize input and make a prediction
26inputs = tokenizer(prompt, return_tensors="pt", max_length=max_input_size, truncation=True)
27model_output = model.generate(
28 inputs["input_ids"],
29 max_new_tokens=100,
30 do_sample=True,
31 temperature=temperature,
32 top_k=top_k,
33 top_p=top_p,
34 num_return_sequences=1,
35 pad_token_id=tokenizer.eos_token_id,
36 attention_mask=inputs["attention_mask"]
37)[0]
38
39# Detokenize result
40output_response = tokenizer.decode(model_output, skip_special_tokens=False)
41
42if "<TERMS>" in output_response:
43 generated_output = output_response.split("<TERMS>")[1].split("<|begin_of_text|>")[0].strip()
44
45 # Print prediction
46 print(generated_output)