feat: change logic

src/streamlit_app.py

@@ -5,35 +5,35 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # from huggingface_hub import notebook_login
 # from unsloth import FastLanguageModel, is_bfloat16_supported
 
-model_id = "sentence-transformers/all-MiniLM-L6-v2"
+# model_id = "sentence-transformers/all-MiniLM-L6-v2"
 # model_id = "sentence-transformers/xlm-r-base-en-ko-nli-ststb"
 
 # model_id = "mistralai/Mistral-7B-Instruct-v0.1"
 # model_id = "meta-llama/Llama-3.2-1B"
-# model_id = "choco-conoz/TwinLlama-3.1-8B"
+model_id = "choco-conoz/TwinLlama-3.1-8B"
 
-# processor = pipeline(
-#     "text-generation",
-#     model=model_id,
-#     model_kwargs={
-#         "torch_dtype": torch.float16,
-#         "quantization_config": {"load_in_4bit": True},
-#         "low_cpu_mem_usage": True,
-#     },
-# )
-# terminators = [
-#     processor.tokenizer.eos_token_id,
-#     processor.tokenizer.convert_tokens_to_ids(""),
-# ]
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
 processor = pipeline(
     "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=10
+    model=model_id,
+    model_kwargs={
+        "torch_dtype": torch.float16,
+        "quantization_config": {"load_in_4bit": True},
+        "low_cpu_mem_usage": True,
+    },
 )
+terminators = [
+    processor.tokenizer.eos_token_id,
+    processor.tokenizer.convert_tokens_to_ids(""),
+]
+
+# tokenizer = AutoTokenizer.from_pretrained(model_id)
+# model = AutoModelForCausalLM.from_pretrained(model_id)
+# processor = pipeline(
+#     "text-generation",
+#     model=model,
+#     tokenizer=tokenizer,
+#     max_new_tokens=10
+# )
 
 
 def main():
@@ -55,11 +55,11 @@ def main():
         prompt = processor.tokenizer.apply_chat_template(
             user_prompt, tokenize=False, add_generation_prompt=True)
         # prompt = user_prompt
-        outputs = processor(prompt)
+        # outputs = processor(prompt)
 
-        # outputs = processor(prompt, max_new_tokens=4096, eos_token_id=terminators, do_sample=True,
-        #                     temperature=0.6, top_p=0.9
-        #                     )
+        outputs = processor(prompt, max_new_tokens=4096, eos_token_id=terminators, do_sample=True,
+                            temperature=0.6, top_p=0.9
+                            )
         response = outputs[0]["generated_text"][len(prompt):]
         st.write(response)