Update app.py

app.py

@@ -3,7 +3,9 @@ import streamlit as st
 from transformers import pipeline
 from PIL import Image
 import os
-
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer
+import torch
 
 def main():
 
@@ -16,14 +18,16 @@ def main():
         return tokens_list
 
 
-    def translate(text, tokenizer, model, bad_words_ids):
+    def translate(text, tokenizer, model, do_sample, max_new_tokens, temperature, top_k, top_p, bad_words_ids):
         # Prepare the prompt
-        messages = f"Translate from Korean to English: {text}"
+        prompts = f"Translate from Korean to English: {text}"
+        messages = [{"role": "user", "content": prompts}]
+        
         input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
         prompt_padded_len = len(input_ids[0])
         
         # Generate the translation
-        gen_tokens = model.generate(input_ids, max_length=max_new_tokens, temperature=temperature, top_k=top_k, top_p=top_p, bad_words_ids = bad_words_ids)
+        gen_tokens = model.generate(input_ids, do_sample = do_sample, max_length=max_new_tokens, temperature=temperature, top_k=top_k, top_p=top_p, bad_words_ids = bad_words_ids)
         gen_tokens = [
           gt[prompt_padded_len:] for gt in gen_tokens
         ]
@@ -46,9 +50,7 @@ def main():
     
     hf_token = os.getenv("HF_ACCESS_TOKEN")
     
-    from peft import AutoPeftModelForCausalLM
-    from transformers import AutoTokenizer
-    import torch
+
     
     # attn_implementation = None
     # USE_FLASH_ATTENTION = False
@@ -72,14 +74,16 @@ def main():
     temperature = st.sidebar.slider("Temperature", value=0.3, min_value=0.0, max_value=1.0, step=0.05)
     top_k = st.sidebar.slider("Top-k", min_value=0, max_value=50, value=0)
     top_p = st.sidebar.slider("Top-p", min_value=0.75, max_value=1.0, step=0.05, value=0.9)
+    do_sample = st.selectbox("do_sample: ",
+                     ['True', 'False'])
 
 
     st.subheader("Enter text to translate")
-    input_text = st.text_area("Text to Translate", value= "Korean text here", height=300)
+    input_text = st.text_area("Text to Translate", value= text_default, height=300)
 
     if st.button("Translate"):
         if input_text:
-            translation = translate(input_text, tokenizer, model, bad_words_ids)
+            translation = translate(input_text, model, do_sample = do_sample, max_new_tokens = max_new_tokens, temperature = temperature, top_k = top_k, top_p = top_p, bad_words_ids = bad_words_ids)
             st.text_area("Translated Text", value=translation, height=300)
         else:
             st.error("Please enter some text to translate.")