Spaces:
Sleeping
Sleeping
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
from accelerate import init_empty_weights, load_checkpoint_and_dispatch, dispatch_model, infer_auto_device_map | |
import streamlit as st | |
from huggingface_hub import login | |
import pandas as pd | |
# Token Secret de Hugging Face | |
huggingface_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"] | |
login(huggingface_token) | |
# Cargar el tokenizador y el modelo | |
model_id = "meta-llama/Llama-3.2-1B" | |
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 | |
# Cargar el modelo con disk_offload | |
with init_empty_weights(): | |
model = AutoModelForCausalLM.from_config(model_id) | |
device_map = infer_auto_device_map(model, max_memory={"disk": "2GiB"}, no_split_module_classes=["LlamaDecoderLayer"]) | |
model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, offload_folder="offload_dir") | |
MAX_INPUT_TOKEN_LENGTH = 10000 | |
def generate_response(input_text, temperature=0.7, max_new_tokens=20): | |
input_ids = tokenizer.encode(input_text, return_tensors='pt').to("cpu") # Usar 'cpu' para mantener la compatibilidad | |
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.") | |
streamer = TextIteratorStreamer(tokenizer, timeout=120.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=20, | |
top_p=0.9, | |
temperature=temperature, | |
num_return_sequences=3, | |
eos_token_id=tokenizer.eos_token_id | |
) | |
try: | |
outputs = model.generate(**generate_kwargs) | |
response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip() | |
return response.split("\n")[0] | |
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") | |
uploaded_file = st.file_uploader("Por favor, sube un archivo CSV para iniciar:", type=["csv"]) | |
if uploaded_file is not None: | |
df = pd.read_csv(uploaded_file) | |
query = 'aspiring human resources specialist' | |
if 'job_title' in df.columns: | |
job_titles = df['job_title'].tolist() | |
# Definir el prompt con in-context learning | |
initial_prompt = ( | |
f"Extract the first record from the dataframe df.\n" | |
f"First job title: '{df.iloc[0]['job_title']}'\n" | |
f"Calculate the cosine similarity between this job title and the query: '{query}'.\n" | |
"Print the cosine similarity score." | |
) | |
st.write("Prompt inicial con In-context Learning:\n") | |
st.write(initial_prompt) | |
if st.button("Generar respuesta"): | |
with st.spinner("Generando respuesta..."): | |
response = generate_response(initial_prompt, temperature=0.5) | |
if response: | |
st.write(f"Respuesta del modelo: {response}") | |
else: | |
st.warning("No se pudo generar una respuesta.") | |
st.success("La conversaci贸n ha terminado.") | |
if st.button("Iniciar nueva conversaci贸n"): | |
st.experimental_rerun() | |
elif st.button("Terminar"): | |
st.stop() | |
else: | |
st.error("La columna 'job_title' no se encuentra en el archivo CSV.") | |
if __name__ == "__main__": | |
main() | |