feat: change pipeline

src/streamlit_app.py

@@ -1,7 +1,7 @@
 import os
 import streamlit as st
 import torch
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # from huggingface_hub import notebook_login
 # from unsloth import FastLanguageModel, is_bfloat16_supported
 
@@ -12,21 +12,29 @@ from transformers import pipeline
 # model_id = "meta-llama/Llama-3.2-1B"
 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_id,
-    model_kwargs={
-        "torch_dtype": torch.float16,
-        "quantization_config": {"load_in_4bit": True},
-        "low_cpu_mem_usage": True,
-    },
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=10
 )
 
-terminators = [
-    processor.tokenizer.eos_token_id,
-    processor.tokenizer.convert_tokens_to_ids(""),
-]
-
 
 def main():
     st.title('Text Generator')
@@ -46,9 +54,12 @@ def main():
         print('user_prompt', user_prompt)
         prompt = processor.tokenizer.apply_chat_template(
             user_prompt, tokenize=False, add_generation_prompt=True)
-        outputs = processor(prompt, max_new_tokens=4096, eos_token_id=terminators, do_sample=True,
-                            temperature=0.6, top_p=0.9
-                            )
+        # prompt = user_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
+        #                     )
         response = outputs[0]["generated_text"][len(prompt):]
         st.write(response)