Upload app.py

app.py

@@ -1,10 +1,11 @@
-from typing import List, Tuple
+from threading import Thread
+from typing import Iterator, List, Tuple
 
 import gradio as gr
 from gradio.themes import Soft
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 TEAM_LOGO_URL = "http://nlp.polytechnique.fr/static/images/logo_dascim.png"
 PROTEIN_VISUAL_URL = "https://cas-bridge.xethub.hf.co/xet-bridge-us/68e677c594d3f20bbeecf13c/7cff6ae021d7c518ee4e2fcb70490516ad9e4999ec75c6a5dd164cc6ca64ae30?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Content-Sha256=UNSIGNED-PAYLOAD&X-Amz-Credential=cas%2F20251023%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20251023T094659Z&X-Amz-Expires=3600&X-Amz-Signature=6a7598d77a46df971e88e1f378bc5e06794a3893f31319a6ab3431e4323d755c&X-Amz-SignedHeaders=host&X-Xet-Cas-Uid=66448b4fecac3bc79b26304f&response-content-disposition=inline%3B+filename*%3DUTF-8%27%27model.png%3B+filename%3D%22model.png%22%3B&response-content-type=image%2Fpng&x-id=GetObject&Expires=1761216419&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTc2MTIxNjQxOX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2FzLWJyaWRnZS54ZXRodWIuaGYuY28veGV0LWJyaWRnZS11cy82OGU2NzdjNTk0ZDNmMjBiYmVlY2YxM2MvN2NmZjZhZTAyMWQ3YzUxOGVlNGUyZmNiNzA0OTA1MTZhZDllNDk5OWVjNzVjNmE1ZGQxNjRjYzZjYTY0YWUzMCoifV19&Signature=YjrX1ZF%7EX1qw-m2nWOY8AxdSXwbrsidvlTZ5YWXZx3UPv0my0u68lWcpWIpIxzkGeWTtWPvlCfMcmnpmmwS2wHexorhgq9c7%7E3Ghw20evO0EMPvHBwP4vWYmXW8nHBqqqbw8Qy1pojDm9TvXV19O4-fCFxPi1aQ5FOTC2Kmn9gKxW%7EAN7vkWnfhU8QcCf18139hMbUvh9YoJ%7EesOWXoCFWgAbyz%7Eroajt5e3oM9b-IsU%7E2-UzMZ4%7EMA2MSOFmg487bhZDbr2IMD15-8O0jzWu3qyO3T1H06S-9kTdI%7EC6AYtXUY8YtSWKw%7EBzhARjXK6%7EuZ3c3kE1V7%7EdnLl1YM-2w__&Key-Pair-Id=K2L8F4GPSG1IFC"
@@ -16,7 +17,7 @@ PROTEIN_HERO = f"""
 """
 
 DESCRIPTION = f"""\
-### Prot2Text-V2 Demo (Spaces Edition)
+### Prot2Text-V2 Demo
 
 {PROTEIN_HERO}
 
@@ -44,17 +45,22 @@ EXAMPLE_SEQUENCES = [
     [
         "MAVVLPAVVEELLSEMAAAVQESARIPDEYLLSLKFLFGSSATQALDLVDRQSITLISSPSGRRVYQVLGSSSKTYTCLASCHYCSCPAFAFSVLRKSDSILCKHLLAVYLSQVMRTCQQLSVSDKQLTDILLMEKKQEA"
     ],
+    [
+        "MCYSANGNTFLIVDNTQKRIPEEKKPDFVRENVGDLDGVIFVELVDGKYFMDYYNRDGSMAAFCGNGARAFSQYLIDRGWIKEKEFTFLSRAGEIKVIVDDSIWVRMPGVSEKKEMKVDGYEGYFVVVGVPHFVMEVKGIDELDVEKLGRDLRYKTGANVDFYEVLPDRLKVRTYERGVERETKACGTGVTSVFVVYRDKTGAKEVKIQVPGGTLFLKEENGEIFLRGDVKRCSEE"
+    ],
+    [
+        "MTQEERFEQRIAQETAIEPQDWMPDAYRKTLIRQIGQHAHSEIVGMLPEGNWITRAPTLRRKAILLAKVQDEAGHGLYLYSAAETLGCAREDIYQKMLDGRMKYSSIFNYPTLSWADIGVIGWLVDGAAIVNQVALCRTSYGPYARAMVKICKEESFHQRQGFEACMALAQGSEAQKQMLQDAINRFWWPALMMFGPNDDNSPNSARSLTWKIKRFTNDELRQRFVDNTVPQVEMLGMTVPDPDLHFDTESGHYRFGEIDWQEFNEVINGRGICNQERLDAKRKAWEEGTWVREAALAHAQKQHARKVA"
+    ],
+    [
+        "MTTRMIILNGGSSAGKSGIVRCLQSVLPEPWLAFGVDSLIEAMPLKMQSAEGGIEFDADGGVSIGPEFRALEGAWAEGVVAMARAGARIIIDDVFLGGAAAQERWRSFVGDLDVLWVGVRCDGAVAEGRETARGDRVAGMAAKQAYVVHEGVEYDVEVDTTHKESIECAWAIAAHVVP"
+    ],
 ]
 
 MAX_MAX_NEW_TOKENS = 256
 DEFAULT_MAX_NEW_TOKENS = 100
 
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-if device.type == "cuda":
-    dtype = torch.bfloat16 if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported() else torch.float16
-else:
-    dtype = torch.float32
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 system_message = (
     "You are a scientific assistant specialized in protein function "
@@ -67,20 +73,17 @@ placeholder = '<|reserved_special_token_1|>'
 
 esm_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
 llama_tokenizer = AutoTokenizer.from_pretrained(
-    pretrained_model_name_or_path="meta-llama/Llama-3.1-8B-Instruct",
-    pad_token='<|reserved_special_token_0|>',
-)
-model = AutoModelForCausalLM.from_pretrained(
-    "xiao-fei/Prot2Text-V2-11B-Instruct-hf",
-    trust_remote_code=True,
-    torch_dtype=dtype,
+    pretrained_model_name_or_path="meta-llama/Llama-3.1-8B-Instruct", 
+    pad_token='<|reserved_special_token_0|>'
 )
-model = model.to(device)
+model = AutoModelForCausalLM.from_pretrained('xiao-fei/Prot2Text-V2-11B-Instruct-hf',
+                                            trust_remote_code=True,
+                                            torch_dtype=torch.bfloat16,).to(device)
 model.eval()
 
 
 @spaces.GPU(duration=90)
-def generate_response(
+def stream_response(
     message: str,
     max_new_tokens: int = 1024,
     do_sample: bool = False,
@@ -88,42 +91,52 @@ def generate_response(
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
-) -> str:
-    user_message = "Sequence embeddings: " + placeholder * (len(message) + 2)
+) -> Iterator[str]:
+
+
+    streamer = TextIteratorStreamer(llama_tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+
+    user_message = "Sequence embeddings: " + placeholder * (len(message)+2)
     tokenized_prompt = llama_tokenizer.apply_chat_template(
         [
             {"role": "system", "content": system_message},
-            {"role": "user", "content": user_message},
-        ],
-        add_generation_prompt=True,
-        tokenize=True,
-        return_tensors="pt",
-        return_dict=True,
+            {"role": "user", "content": user_message}
+        ], 
+         add_generation_prompt=True, 
+        tokenize=True, 
+        return_tensors="pt", 
+        return_dict=True
     )
-    tokenized_sequence = esm_tokenizer(message, return_tensors="pt")
-
-    with torch.inference_mode():
-        generated = model.generate(
-            inputs=tokenized_prompt["input_ids"].to(device),
-            attention_mask=tokenized_prompt["attention_mask"].to(device),
-            protein_input_ids=tokenized_sequence["input_ids"].to(device),
-            protein_attention_mask=tokenized_sequence["attention_mask"].to(device),
-            eos_token_id=128009,
-            pad_token_id=128002,
-            return_dict_in_generate=False,
-            num_beams=1,
-            max_new_tokens=max_new_tokens,
-            do_sample=do_sample,
-            top_p=top_p,
-            top_k=top_k,
-            temperature=temperature,
-            repetition_penalty=repetition_penalty,
-        )
+    tokenized_sequence = esm_tokenizer(
+        message, 
+        return_tensors="pt"
+    )
+    model.eval()    
+    generate_kwargs = dict(
+        inputs=tokenized_prompt["input_ids"].to(model.device),
+        attention_mask=tokenized_prompt["attention_mask"].to(model.device),
+        protein_input_ids=tokenized_sequence["input_ids"].to(model.device),
+        protein_attention_mask=tokenized_sequence["attention_mask"].to(model.device),
+        eos_token_id=128009, 
+        pad_token_id=128002,
+        return_dict_in_generate=False,
+        num_beams=1,
+        # device=device,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
 
-    prompt_len = tokenized_prompt["input_ids"].shape[-1]
-    response_tokens = generated[0, prompt_len:]
-    response_text = llama_tokenizer.decode(response_tokens, skip_special_tokens=True)
-    return response_text.strip()
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
 
 
 ChatHistory = List[Tuple[str, str]]
@@ -147,30 +160,29 @@ def handle_submit(
     conversation = history.copy()
     conversation.append((message, ""))
 
-    try:
-        response_text = generate_response(
-            message=message,
-            max_new_tokens=max_new_tokens,
-            do_sample=do_sample,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-        )
-    except Exception as exc:
-        response_text = f"⚠️ Generation failed: {exc}"
-
-    conversation[-1] = (message, response_text)
-    snapshot = conversation.copy()
-    return snapshot, snapshot, gr.update(value="")
+    for partial_response in stream_response(
+        message=message,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+    ):
+        conversation[-1] = (message, partial_response)
+        snapshot = conversation.copy()
+        yield snapshot, snapshot, gr.update(value="")
 
 
 def clear_conversation():
     empty_history: ChatHistory = []
     return empty_history, empty_history, gr.update(value="")
 
-
-theme = Soft(primary_hue="slate", secondary_hue="stone", neutral_hue="gray")
+theme = Soft(
+    primary_hue="slate",
+    secondary_hue="stone",
+    neutral_hue="gray",
+)
 
 with gr.Blocks(theme=theme, css_paths="style.css", fill_height=True) as demo:
     with gr.Row(equal_height=True):
@@ -187,6 +199,8 @@ with gr.Blocks(theme=theme, css_paths="style.css", fill_height=True) as demo:
             gr.Markdown(DESCRIPTION)
         with gr.Column(scale=7, min_width=400, elem_classes="interaction-column"):
             history_state = gr.State([])
+            
+
             chatbot = gr.Chatbot(
                 label="Generated Function",
                 height=350,
@@ -289,10 +303,9 @@ with gr.Blocks(theme=theme, css_paths="style.css", fill_height=True) as demo:
             )
     with gr.Accordion("Model & usage notes", open=False):
         gr.Markdown(
-            "- **Model stack**: Facebook ESM2 encoder + Prot2Text-V2 11B instruction-tuned decoder.\n"
+            "- **Model stack**: Facebook ESM2 encoder + Llama 3.1 8B instruction-tuned decoder.\n"
             "- **Token budget**: the generator truncates after the configured `Max new tokens`.\n"
             "- **Attribution**: Outputs are predictions; validate experimentally before publication.\n"
-            "- **Privacy**: Sequence inputs stay within this session - export or clear as needed."
         )
     gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")