feat: 1

src/streamlit_app.py

@@ -12,29 +12,30 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, 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,
+#     },
+# )
+
+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(""),
+    tokenizer.eos_token_id,
+    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():
     st.title('Text Generator')