Update app.py

app.py

@@ -14,10 +14,10 @@ login(huggingface_token)
 # model_name = "meta-llama/Llama-3.2-3B-Instruct"
 model_name = "meta-llama/Llama-3.2-1B-Instruct"
 model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name, device_map = 'auto')
 
 #pipe = pipeline(model=model, tokenizer=tokenizer, max_length = 512) # Check documentation without "feature-extraction"
-pipe = pipeline(task = 'text-generation', model=model, tokenizer=tokenizer, max_length = 512) # Check documentation without "feature-extraction"
+pipe = pipeline(task = 'text-generation', model=model, tokenizer=tokenizer, max_length = 512, device = 'auto') # Check documentation without "feature-extraction"
 
 
 # Use transformers pipeline
@@ -35,47 +35,38 @@ uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
 print("Query: ", query)
 
 if uploaded_file is not None:
-    # Read CSV file
     df = pd.read_csv(uploaded_file)
 
-    job_title = df['job_title']
-    # Request a query to user
-    #query = st.text_input("Enter your query:")
-    
-    if query:
-        st.write("Query", query)
-        st.write("Job Title:" , job_title)
-        # Create the prompt to calculate cosine similarity score
-        prompt = f"""
-        I will provide you a list of candidates' job titles and the search query, I need you to calculate the cosine similarity score between each record and the query.
-        You should create a new column called 'Score' with those scores and then sort the records from highest to lowest 'Score'.
-        The query is: {query} and the list of job titles is: {job_title}
-        """
-
-        # Call the model with the prompt
-        try:
-            response = llm_pipeline(prompt, max_length = 512, max_new_tokens = 300, truncation = True, padding = True)
-            # Show Model answer
-            st.write("Model Answer:", response)
-
-        except Exception as e:
-            st.error(f"Error while processing: {str(e)}")
-
-# Data validation Pydantic
-class ConsultModel(BaseModel):
-    query: str
-
-    @model_validator(mode='before')
-    def validate_query(cls, values):
-        query = values.get('query')
-        if not query:
-            raise ValueError("Query cannot be empty.")
-        return values
-
-# Example
-if query:
-    try:
-        valid_query = ConsultModel(query=query)
-        st.success("Valid consult.")
-    except ValueError as e:
-        st.error(f"Validation error: {e}")
+    if 'job_title' not in df.columns:
+        st.error("The uploaded CSV must contain a 'job_title' column.")
+    else:
+        job_titles = df['job_title'].tolist()
+
+        if query:
+            st.write("Query:", query)
+
+            # Crear el prompt para el LLM
+            prompt = f"""
+            You are an AI assistant. You will be given a list of job titles and a search query. 
+            Your task is to calculate the cosine similarity score between the query and each job title.
+            You should then return a sorted list of job titles based on the similarity score, from highest to lowest.
+            Provide the output in the following format:
+
+            1. Job Title: [Job Title], Score: [Cosine Similarity Score]
+            2. Job Title: [Job Title], Score: [Cosine Similarity Score]
+            ...
+            
+            Query: {query}
+            Job Titles: {job_titles}
+            """
+
+            # Llamar al modelo con el prompt
+            try:
+                response = pipe(prompt, max_length=1024, num_return_sequences=1)
+
+                # Mostrar la respuesta del modelo
+                st.write("Model Answer:")
+                st.write(response[0]['generated_text'])
+
+            except Exception as e:
+                st.error(f"Error while processing: {str(e)}")