Upload 2 files

app.py

@@ -1,11 +1,10 @@
-from threading import Thread
-from typing import Iterator, List, Tuple
+from typing import List, Tuple
 
 import gradio as gr
 from gradio.themes import Soft
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 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"
@@ -17,7 +16,7 @@ PROTEIN_HERO = f"""
 """
 
 DESCRIPTION = f"""\
-### Prot2Text-V2 Demo
+### Prot2Text-V2 Demo (Spaces Edition)
 
 {PROTEIN_HERO}
 
@@ -45,22 +44,17 @@ EXAMPLE_SEQUENCES = [
     [
         "MAVVLPAVVEELLSEMAAAVQESARIPDEYLLSLKFLFGSSATQALDLVDRQSITLISSPSGRRVYQVLGSSSKTYTCLASCHYCSCPAFAFSVLRKSDSILCKHLLAVYLSQVMRTCQQLSVSDKQLTDILLMEKKQEA"
     ],
-    [
-        "MCYSANGNTFLIVDNTQKRIPEEKKPDFVRENVGDLDGVIFVELVDGKYFMDYYNRDGSMAAFCGNGARAFSQYLIDRGWIKEKEFTFLSRAGEIKVIVDDSIWVRMPGVSEKKEMKVDGYEGYFVVVGVPHFVMEVKGIDELDVEKLGRDLRYKTGANVDFYEVLPDRLKVRTYERGVERETKACGTGVTSVFVVYRDKTGAKEVKIQVPGGTLFLKEENGEIFLRGDVKRCSEE"
-    ],
-    [
-        "MTQEERFEQRIAQETAIEPQDWMPDAYRKTLIRQIGQHAHSEIVGMLPEGNWITRAPTLRRKAILLAKVQDEAGHGLYLYSAAETLGCAREDIYQKMLDGRMKYSSIFNYPTLSWADIGVIGWLVDGAAIVNQVALCRTSYGPYARAMVKICKEESFHQRQGFEACMALAQGSEAQKQMLQDAINRFWWPALMMFGPNDDNSPNSARSLTWKIKRFTNDELRQRFVDNTVPQVEMLGMTVPDPDLHFDTESGHYRFGEIDWQEFNEVINGRGICNQERLDAKRKAWEEGTWVREAALAHAQKQHARKVA"
-    ],
-    [
-        "MTTRMIILNGGSSAGKSGIVRCLQSVLPEPWLAFGVDSLIEAMPLKMQSAEGGIEFDADGGVSIGPEFRALEGAWAEGVVAMARAGARIIIDDVFLGGAAAQERWRSFVGDLDVLWVGVRCDGAVAEGRETARGDRVAGMAAKQAYVVHEGVEYDVEVDTTHKESIECAWAIAAHVVP"
-    ],
 ]
 
 MAX_MAX_NEW_TOKENS = 256
 DEFAULT_MAX_NEW_TOKENS = 100
 
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+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
 
 system_message = (
     "You are a scientific assistant specialized in protein function "
@@ -73,17 +67,20 @@ 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|>'
+    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,
 )
-model = AutoModelForCausalLM.from_pretrained('xiao-fei/Prot2Text-V2-11B-Instruct-hf',
-                                            trust_remote_code=True,
-                                            torch_dtype=torch.bfloat16,).to(device)
+model = model.to(device)
 model.eval()
 
 
 @spaces.GPU(duration=90)
-def stream_response(
+def generate_response(
     message: str,
     max_new_tokens: int = 1024,
     do_sample: bool = False,
@@ -91,52 +88,42 @@ def stream_response(
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.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)
+) -> str:
+    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
-    )
-    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,
+            {"role": "user", "content": user_message},
+        ],
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt",
+        return_dict=True,
     )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+    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,
+        )
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    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()
 
 
 ChatHistory = List[Tuple[str, str]]
@@ -160,29 +147,30 @@ def handle_submit(
     conversation = history.copy()
     conversation.append((message, ""))
 
-    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="")
+    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()
+    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):
@@ -199,8 +187,6 @@ 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,
@@ -303,9 +289,10 @@ 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 + Llama 3.1 8B instruction-tuned decoder.\n"
+            "- **Model stack**: Facebook ESM2 encoder + Prot2Text-V2 11B 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")
 

requirements.txt

@@ -4,5 +4,6 @@ gradio>=5.0.0
 huggingface-hub>=0.23.0
 mistral-common>=1.4.0
 sentencepiece>=0.1.99
+spaces>=0.27.0
 torch>=2.1.0
 transformers>=4.38.0