Upload app.py

app.py

@@ -6,23 +6,13 @@ import os
 Copied from inference in colab notebook
 """
 
-from transformers import AutoTokenizer , AutoModelForSeq2SeqLM , TextIteratorStreamer
+from transformers import TextIteratorStreamer , pipeline
 from threading import Thread
 
 # 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)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_path, token = os.getenv('huggingface_token'))
-
-# # Load the base model (e.g., LLaMA)
-# base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model, token = os.getenv('huggingface_token'))
-
-# # Load LoRA adapter
-# from peft import PeftModel
-# model = PeftModel.from_pretrained(base_model, model_path, token = os.getenv('huggingface_token'))
+translator = pipeline("translation_xx_to_yy", model=model_path)
 
 def respond(
     message: str,
@@ -32,40 +22,54 @@ 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})
+    message = "translate English to French:" + message
+
+    response = translator(message)
+    print(response)
+    yield response
+
+# def respond(
+#     message: str,
+#     history: list[tuple[str, str]],
+#     system_message: str,
+#     max_tokens: int,
+#     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()
+#     # 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
+#     # Yield responses as they are generated
+#     response = ""
+#     for token in streamer:
+#         response += token
+#         yield response
 
 
 """