Update app.py

app.py

@@ -6,86 +6,154 @@ import os
 Copied from inference in colab notebook
 """
 
-from transformers import AutoTokenizer , AutoModelForSeq2SeqLM , TextIteratorStreamer
-from threading import Thread
+from transformers import pipeline
 
 # Load model and tokenizer globally to avoid reloading for every request
-base_model = "google-t5/t5-small"
 model_path = "Mat17892/t5small_enfr_opus"
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, legacy=False)
+# translator = pipeline("translation_xx_to_yy", model=model_path)
 
-# Load the base model (e.g., LLaMA)
-base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model, token = os.getenv('huggingface_token'))
+# def respond(
+#     message: str,
+#     history: list[tuple[str, str]],
+#     system_message: str,
+#     max_tokens: int,
+#     temperature: float,
+#     top_p: float,
+# ):
+#     message = "translate English to French:" + message
 
-# Load LoRA adapter
-from peft import PeftModel
-model = PeftModel.from_pretrained(base_model, model_path, token = os.getenv('huggingface_token'))
+#     response = translator(message)[0]
+#     yield response['translation_text']
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextIteratorStreamer
+import threading
+
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
 
 def respond(
     message: str,
     history: list[tuple[str, str]],
     system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
+    max_tokens: int = 128,
+    temperature: float = 1.0,
+    top_p: float = 1.0,
 ):
-    # Combine system message and history into a single prompt
-    messages = [{"role": "system", "content": system_message}]
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-    messages.append({"role": "user", "content": message})
-    
-    # Tokenize the messages
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        tokenize = True,
-        add_generation_prompt = True, # Must add for generation
-        return_tensors = "pt",
+    # Preprocess the input message
+    input_text = system_message + " " + message
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+
+    # Set up the streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+
+    # Generate in a separate thread to avoid blocking
+    generation_thread = threading.Thread(
+        target=model.generate,
+        kwargs={
+            "input_ids": input_ids,
+            "max_new_tokens": max_tokens,
+            "do_sample": True,
+            "temperature": temperature,
+            "top_p": top_p,
+            "streamer": streamer,
+        },
     )
-    # Generate tokens incrementally
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {
-        "input_ids": inputs,
-        "max_new_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": top_p,
-        "do_sample": True,
-        "streamer": streamer,
-    }
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    # Yield responses as they are generated
-    response = ""
-    for token in streamer:
-        response += token
-        yield response
+    generation_thread.start()
+
+    # Stream the output progressively
+    for token in streamer:  # Append each token to the accumulated text
+        yield token
 
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+# demo = gr.ChatInterface(
+#     respond,
+#     additional_inputs=[
+#         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+#         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+#         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+#         gr.Slider(
+#             minimum=0.1,
+#             maximum=1.0,
+#             value=0.95,
+#             step=0.05,
+#             label="Top-p (nucleus sampling)",
+#         ),
+#     ],
+# )
+
+# Function to process translation
+def respond_google_translate(
+    source_text,
+    system_message,
+    max_tokens,
+    temperature,
+    top_p
+):
+    # Call the respond function and collect the final response
+    result = ""
+    for token in respond(
+        message=source_text,
+        history=[],
+        system_message=system_message,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        result += token  # Accumulate the tokens
+    return result
+
+# Define the interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Google Translate-like Interface")
+
+    with gr.Row():
+        with gr.Column():
+            source_textbox = gr.Textbox(
+                placeholder="Enter text in English...",
+                label="Source Text (English)",
+                lines=5,
+            )
+        with gr.Column():
+            translated_textbox = gr.Textbox(
+                placeholder="Translation will appear here...",
+                label="Translated Text (French)",
+                lines=5,
+                interactive=False,
+            )
 
+    translate_button = gr.Button("Translate")
+
+    with gr.Accordion("Advanced Settings", open=False):
+        system_message_input = gr.Textbox(
+            value="translate English to French:",
+            label="System message",
+        )
+        max_tokens_slider = gr.Slider(
+            minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"
+        )
+        temperature_slider = gr.Slider(
+            minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"
+        )
+        top_p_slider = gr.Slider(
+            minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"
+        )
+
+    # Define functionality
+    translate_button.click(
+        respond_google_translate,
+        inputs=[
+            source_textbox,
+            system_message_input,
+            max_tokens_slider,
+            temperature_slider,
+            top_p_slider,
+        ],
+        outputs=translated_textbox,
+    )
 
 if __name__ == "__main__":
     demo.launch()