Upload app.py

app.py

@@ -6,27 +6,12 @@ import os
 Copied from inference in colab notebook
 """
 
-from transformers import TextIteratorStreamer , pipeline
-from threading import Thread
+from transformers import pipeline
 
 # Load model and tokenizer globally to avoid reloading for every request
 model_path = "Mat17892/t5small_enfr_opus"
 
-translator = pipeline("translation_xx_to_yy", model=model_path)
-
-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
-
-    response = translator(message)[0]
-    print(response)
-    yield response['translation_text']
+# translator = pipeline("translation_xx_to_yy", model=model_path)
 
 # def respond(
 #     message: str,
@@ -36,40 +21,49 @@ def respond(
 #     temperature: float,
 #     top_p: float,
 # ):
-#     # 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",
-#     )
-#     # 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
+#     message = "translate English to French:" + message
+
+#     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 = 128,
+    temperature: float = 1.0,
+    top_p: float = 1.0,
+):
+    # Preprocess the input message
+    input_text = "translate English to French: " + 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,
+        },
+    )
+    generation_thread.start()
+
+    # Stream the output progressively
+    for token in streamer:
+        yield token
 
 
 """