Update app.py

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import spaces
 import selfies as sf
 from llama_cpp import Llama
 from llama_cpp_agent import LlamaCppAgent
@@ -26,16 +25,7 @@ css = """
 }
 """
 
-@spaces.GPU(duration=120)
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    max_tokens,
-    temperature,
-    top_p,
-    top_k
-):
-
+def respond(encoded_smiles, max_tokens, temperature, top_p, top_k):
     model_name = "model.gguf"
     llm = Llama(model_name)
     provider = LlamaCppPythonProvider(llm)
@@ -52,46 +42,23 @@ def respond(
     settings.top_k = top_k
     settings.top_p = top_p
     settings.max_tokens = max_tokens
-    settings.stream = True
+    settings.stream = False
 
-    messages = BasicChatHistory()
-    
-    for msn in history:
-        user = {
-            'role': Roles.user,
-            'content': msn[0]
-        }
-        assistant = {
-            'role': Roles.assistant,
-            'content': msn[1]
-        }
-        messages.add_message(user)
-        messages.add_message(assistant)
+    prompt = f"{encoded_smiles}"
+    input_ids = agent.tokenizer(prompt, return_tensors='pt', truncation=False).input_ids.cuda()
+    outputs = agent.model.generate(input_ids=input_ids, max_new_tokens=512)
+    output1 = agent.tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]
 
-    stream = agent.get_chat_response(
-        sf.encoder(message),
-        llm_sampling_settings=settings,
-        chat_history=messages,
-        returns_streaming_generator=True,
-        print_output=False
-    )
+    first_inst_index = output1.find("[/INST]")
+    second_inst_index = output1.find("[/IN", first_inst_index + len("[/INST]") + 1)
+    predicted_selfies = output1[first_inst_index + len("[/INST]") : second_inst_index].strip()
 
-    outputs = ""
-    for output in stream:
-        outputs += output
-        yield outputs
-
-PLACEHOLDER = """
-<div class="message-bubble-border" style="display:flex; max-width: 600px; border-radius: 8px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); backdrop-filter: blur(10px);">
-    <div style="padding: .5rem 1.5rem;">
-        <h2 style="text-align: left; font-size: 1.5rem; font-weight: 800; margin-bottom: 0.5rem;">Retrosynthesis Chatbot</h2>
-    </div>
-</div>
-"""
+    return {'input': encoded_smiles, 'predict': predicted_selfies}
 
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
+demo = gr.Interface(
+    fn=respond,
+    inputs=[
+        gr.Textbox(label="Encoded SMILES"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=1.0, step=0.1, label="Temperature"),
         gr.Slider(
@@ -109,6 +76,7 @@ demo = gr.ChatInterface(
             label="Top-k",
         )
     ],
+    outputs=gr.JSON(label="Results"),
     theme=gr.themes.Soft(primary_hue="violet", secondary_hue="violet", neutral_hue="gray", font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]).set(
         body_background_fill_dark="#16141c",
         block_background_fill_dark="#16141c",
@@ -121,13 +89,8 @@ demo = gr.ChatInterface(
         color_accent_soft_dark="transparent"
     ),
     css=css,
-    retry_btn="Retry",
-    undo_btn="Undo",
-    clear_btn="Clear",
-    submit_btn="Send",
     description="Retrosynthesis chatbot",
-    chatbot=gr.Chatbot(scale=1, placeholder=PLACEHOLDER)
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()