Update app.py

app.py

@@ -35,12 +35,13 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
 
 
 def run(message: str,
-        chat_history: list[tuple[str, str]],
-        system_prompt: str,
-        max_new_tokens: int = 1024,
-        temperature: float = 0.1,
-        top_p: float = 0.9,
-        top_k: int = 50) -> Iterator[str]:
+                   chat_history: list[tuple[str, str]],
+                   system_prompt: str,
+                   max_new_tokens: int = 1024,
+                   temperature: float = 0.1,
+                   top_p: float = 0.9,
+                   top_k: int = 50) -> Iterator[str]:
+
     prompt = get_prompt(message, chat_history, system_prompt)
 
     generate_kwargs = dict(
@@ -50,23 +51,17 @@ def run(message: str,
         top_k=top_k,
         temperature=temperature,
     )
-    -> Iterator[list[tuple[str, str]]]:
-    if max_new_tokens > MAX_MAX_NEW_TOKENS:
-        raise ValueError
 
-    history = history_with_input[:-1]
-    generator = run(message, history, system_prompt, max_new_tokens, temperature, top_p, top_k)
-    
-    try:
-        first_response = next(generator)
-        translated_response = translator_to_ar.translate(first_response)  # Translate the response to Arabic
-        yield history + [(message, translated_response)]
-    except StopIteration:
-        yield history + [(message, '')]
+    stream = client.generate_stream(prompt, **generate_kwargs)
+    output = ""
     
-    for response in generator:
-        translated_response = translator_to_ar.translate(response)  # Translate each response to Arabic
-        yield history + [(message, translated_response)]
+    for response in stream:
+        if any([end_token in response.token.text for end_token in [EOS_STRING, EOT_STRING]]):
+            translated_output = translator_to_ar.translate(output)
+            yield translated_output
+            output = ""
+        else:
+            output += response.token.text
 
 
 def generate_image_caption(image_data):