Update app.py

app.py

@@ -14,40 +14,43 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
 tokenizer.pad_token = tokenizer.eos_token
 
-MAX_INPUT_TOKEN_LENGTH = 10000
+MAX_INPUT_TOKEN_LENGTH = 4096
 
-# Asegurar que el token de padding esté configurado
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token  # Asignar el token de padding al token de fin de oración
+def generate_response(input_text, temperature=0.5, max_new_tokens=50):
+    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
 
-def generate_response(input_text, max_new_tokens=50, temperature=0.5):
-    """
-    Genera una respuesta usando el modelo de lenguaje con máscara de atención.
-    """
-    # Tokenizar la entrada y crear la máscara de atención
-    inputs = tokenizer(
-        input_text,
-        return_tensors='pt',
-        padding=True,
-        truncation=True,
-        max_length=512  # Ajustar según sea necesario
-    )
-
-    input_ids = inputs['input_ids'].to(model.device)
-    attention_mask = inputs['attention_mask'].to(model.device)
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        st.warning(f"Se recortó la entrada porque excedió el límite de {MAX_INPUT_TOKEN_LENGTH} tokens.")
 
-    # Generar texto con la máscara de atención y el token de padding
-    outputs = model.generate(
-        input_ids,
-        attention_mask=attention_mask,
+    streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
         max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_k=40,
+        top_p=0.9,
         temperature=temperature,
-        pad_token_id=tokenizer.pad_token_id  # Usar el token de padding configurado
+        eos_token_id=[tokenizer.eos_token_id]
     )
 
-    # Decodificar la respuesta generada
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    try:
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+        t.join()  # Asegura que la generación haya terminado
+
+        outputs = []
+        for text in streamer:
+            outputs.append(text)
+        if not outputs:
+            raise ValueError("No se generó ninguna respuesta.")
+        
+        response = "".join(outputs).strip().split("\n")[0]
+        return response
+    except Exception as e:
+        st.error(f"Error durante la generación: {e}")
+        return "Error en la generación de texto."
 
 def main():
     st.title("Chat con Meta Llama 3.2 1B")
@@ -58,23 +61,25 @@ def main():
         df = pd.read_csv(uploaded_file)
         
         if 'job_title' in df.columns:
-            query = "aspiring human resources specialist"
             job_titles = df['job_title'].tolist()
 
             # Definir el prompt con in-context learning
             initial_prompt = (
-                "You are an AI assistant specialized in job title extraction and similarity calculation.\n"
-                "Your task is to extract the first job title from the following list and calculate the cosine similarity with the given query.\n"
+                "Here are some examples of job title extraction:\n"
+                "Example 1:\n"
+                "List: ['Data Scientist', 'Machine Learning Engineer', 'AI Researcher']\n"
+                "First job title: 'Data Scientist'\n"
+                "\n"
+                "Example 2:\n"
+                "List: ['Software Developer', 'Backend Engineer', 'Frontend Developer']\n"
+                "First job title: 'Software Developer'\n"
+                "\n"
+                "Now, extract the first job title from the following list:\n"
                 f"List: {job_titles}\n"
-                f"Query: '{query}'\n"
-                "Extracted first job title: \n"
-                "Cosine similarity score: "
+                "First job title:"
             )
-
-
             
             st.write("Prompt inicial con In-context Learning:")
-            st.write(query)
             st.write(initial_prompt)
 
             if st.button("Generar respuesta"):
@@ -96,4 +101,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-