Update app.py

app.py

@@ -1,28 +1,77 @@
-import os
-import torch
-from transformers import pipeline
 import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_id = "meta-llama/Llama-3.2-1B-Instruct"
+# Cargar modelo
+model_name = "Qwen/Qwen3-1.7B"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto",
+    trust_remote_code=True
+)
 
-st.title("Bienvenido a la app de prueba con Llama Instruct")
-st.write("Prueba basica para la prueba de Huggins Face con el modelo de llama 3.2 instrutc")
+# Cargar prompt base
+with open("SeguroCargo.txt", "r", encoding="utf-8") as f:
+    prompt_base = f.read().strip()
 
-input = st.text_input("Agregue una pregunta")
+# Función para generar respuesta con "thinking" opcional
+def generar_respuesta(pregunta):
+    contexto = f"{prompt_base}\nUsuario: {pregunta}\nAsistente:"
+    messages = [{"role": "user", "content": contexto}]
 
-pipe = pipeline(
-    "text-generation",
-    model=model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    token = os.getenv("HF_TOKEN2")
-)
-messages = [
-    {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
-    {"role": "user", "content": "Who are you?"},
-]
-outputs = pipe(
-    messages,
-    max_new_tokens=256,
-)
-st.write(print(outputs[0]["generated_text"][-1]))
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=False
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=1024,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.8
+    )
+
+    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+
+    try:
+        index = len(output_ids) - output_ids[::-1].index(151668)
+    except ValueError:
+        index = 0
+
+    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+    return content.strip()
+
+# Configuración de Streamlit
+st.set_page_config(page_title="Asistente SeguroCargo2", page_icon="📦")
+st.title("📦 Bienvenido a SeguroCargo2")
+st.markdown("Hola, yo soy **SEGU**, asistente virtual de SeguroCargo sin el modo de pensamiento!")
+st.markdown("Consulta sobre tus **envíos nacionales e internacionales** de manera rápida y profesional.")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Mostrar historial de mensajes
+for msg in st.session_state.messages:
+    role, text = msg
+    with st.chat_message(role):
+        st.markdown(text)
+
+# Entrada del usuario
+user_input = st.chat_input("Escribe tu mensaje...")
+
+if user_input:
+    st.session_state.messages.append(("user", user_input))
+    with st.chat_message("user"):
+        st.markdown(user_input)
+
+    # Generar y mostrar respuesta
+    respuesta = generar_respuesta(user_input)
+    st.session_state.messages.append(("assistant", respuesta))
+    with st.chat_message("assistant"):
+        st.markdown(respuesta)