Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| import streamlit as st | |
| from dotenv import load_dotenv # <-- Added to load your .env file securely | |
| from transformers import pipeline, AutoTokenizer, BitsAndBytesConfig | |
| from peft import AutoPeftModelForCausalLM | |
| # 1. Load environment variables from .env file BEFORE doing anything else | |
| load_dotenv() | |
| # ---------------- CONFIG ---------------- | |
| st.set_page_config(page_title="Kriolu AI Hub", layout="wide") | |
| # Using the token from environment variables | |
| token = os.environ.get("HF_TOKEN") | |
| # ---------------- CACHE (STANDARD) ---------------- | |
| def load_pipeline(task, model_path, **kwargs): | |
| """ | |
| Standard pipeline loader. | |
| Added 4-bit support specifically for the 70B model to prevent OOM. | |
| """ | |
| if "70b" in model_path.lower(): | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_compute_dtype=torch.bfloat16, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_use_double_quant=True, | |
| ) | |
| return pipeline( | |
| task, | |
| model=model_path, | |
| tokenizer=model_path, | |
| token=token, | |
| model_kwargs={"quantization_config": bnb_config, "device_map": "auto"}, | |
| **kwargs | |
| ) | |
| return pipeline(task, model=model_path, tokenizer=model_path, token=token, **kwargs) | |
| # ---------------- CACHE (PEFT / LORA NATIVE) ---------------- | |
| def load_peft_model(model_path): | |
| """Loads 70B model and tokenizer natively using 4-bit quantization.""" | |
| tokenizer = AutoTokenizer.from_pretrained(model_path, token=token) | |
| # 4-bit config is mandatory for 70B on a single 80GB GPU | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_compute_dtype=torch.bfloat16, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_use_double_quant=True, | |
| ) | |
| model = AutoPeftModelForCausalLM.from_pretrained( | |
| model_path, | |
| token=token, | |
| quantization_config=bnb_config, | |
| device_map="auto", | |
| trust_remote_code=True | |
| ) | |
| return model, tokenizer | |
| # ---------------- DECODER (PEFT / LORA NATIVE) ---------------- | |
| def instantiate_peft_decoder(model_name, max_length_, num_return_sequences, text): | |
| model_path = f'Iscte-Sintra/{model_name}' | |
| model, tokenizer = load_peft_model(model_path) | |
| inputs = tokenizer(str(text), return_tensors="pt").to(model.device) | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=max_length_, | |
| num_return_sequences=num_return_sequences, | |
| do_sample=True, | |
| top_p=0.95, | |
| top_k=50, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| results = [] | |
| for output in outputs: | |
| generated_text = tokenizer.decode(output[inputs.input_ids.shape[-1]:], skip_special_tokens=True) | |
| results.append({"generated_text": generated_text}) | |
| return results | |
| # ----------- CHATBOT GENERATION ROUTER ----------- | |
| def generate_chatbot_response(model_name, is_peft, conversation_history): | |
| """Handles chatbot string processing for both PEFT and Standard pipelines.""" | |
| model_path = f"Iscte-Sintra/{model_name}" | |
| if is_peft: | |
| model, tokenizer = load_peft_model(model_path) | |
| else: | |
| pipe = load_pipeline("text-generation", model_path) | |
| if pipe is None: | |
| return "Falha ao inicializar o pipeline." | |
| model = pipe.model | |
| tokenizer = pipe.tokenizer | |
| if getattr(tokenizer, "chat_template", None) is None: | |
| tokenizer.chat_template = ( | |
| "{% for message in messages %}" | |
| "{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n'}}" | |
| "{% endfor %}" | |
| "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}" | |
| ) | |
| prompt_str = tokenizer.apply_chat_template(conversation_history, tokenize=False, add_generation_prompt=True) | |
| if isinstance(prompt_str, list): | |
| prompt_str = prompt_str[0] | |
| pad_id = tokenizer.eos_token_id | |
| if isinstance(pad_id, list): pad_id = pad_id[0] | |
| if pad_id is None: pad_id = 50256 | |
| if is_peft: | |
| inputs = tokenizer(str(prompt_str), return_tensors="pt").to(model.device) | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=150, | |
| do_sample=True, | |
| temperature=0.7, | |
| pad_token_id=pad_id | |
| ) | |
| generated_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True) | |
| else: | |
| response = pipe( | |
| str(prompt_str), | |
| max_new_tokens=150, | |
| do_sample=True, | |
| temperature=0.7, | |
| pad_token_id=pad_id, | |
| return_full_text=False | |
| ) | |
| if isinstance(response, list) and len(response) > 0: | |
| if isinstance(response[0], list): | |
| generated_text = response[0][0].get('generated_text', '') | |
| else: | |
| generated_text = response[0].get('generated_text', '') | |
| else: | |
| generated_text = str(response) | |
| return generated_text.strip() | |
| # ---------------- CHATBOT UI ---------------- | |
| def build_chatbot_page(model_name, is_peft=False): | |
| st.title(f"🤖 {model_name}: Chatbot") | |
| if "current_chat_model" not in st.session_state or st.session_state.current_chat_model != model_name: | |
| st.session_state.messages = [] | |
| st.session_state.current_chat_model = model_name | |
| for message in st.session_state.messages: | |
| with st.chat_message(message["role"]): | |
| st.markdown(message["content"]) | |
| if prompt := st.chat_input("Diz algo..."): | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| with st.chat_message("user"): | |
| st.markdown(prompt) | |
| with st.chat_message("assistant"): | |
| with st.spinner("A pensar..."): | |
| try: | |
| conversation = [{"role": "system", "content": "You are a helpful assistant."}] + st.session_state.messages | |
| generated_text = generate_chatbot_response(model_name, is_peft, conversation) | |
| st.markdown(generated_text) | |
| st.session_state.messages.append({"role": "assistant", "content": generated_text}) | |
| except Exception as e: | |
| st.error(f"Erro na geração: {e}") | |
| # ---------------- ENCODER ---------------- | |
| def instantiate_encoder(model_name, top_k, text): | |
| pipe = load_pipeline("fill-mask", f"Iscte-Sintra/{model_name}") | |
| return pipe(text, top_k=top_k) | |
| # ---------------- TRANSLATION ---------------- | |
| def instantiate_translation_model(model_name, text, src_lg, tgt_lg): | |
| model_path = f'Iscte-Sintra/{model_name}' | |
| if "nllb" in model_name: | |
| pipe = pipeline( | |
| "translation", | |
| model=model_path, | |
| tokenizer=model_path, | |
| token=token, | |
| src_lang=src_lg, | |
| tgt_lang=tgt_lg | |
| ) | |
| return pipe(text)[0]["translation_text"] | |
| elif "m2m100" in model_name: | |
| pipe = load_pipeline("translation", model_path) | |
| pipe.tokenizer.src_lang = src_lg | |
| tgt_lang_id = pipe.tokenizer.convert_tokens_to_ids(tgt_lg) | |
| if tgt_lang_id == pipe.tokenizer.unk_token_id: | |
| st.error(f"Erro: O token {tgt_lg} não foi encontrado no vocabulário do modelo!") | |
| return None | |
| result = pipe(text, forced_bos_token_id=tgt_lang_id) | |
| return result[0]["translation_text"] | |
| else: | |
| pipe = pipeline( | |
| "translation", | |
| model=model_path, | |
| tokenizer=model_path, | |
| token=token, | |
| src_lang=src_lg, | |
| tgt_lang=tgt_lg | |
| ) | |
| return pipe(text)[0]["translation_text"] | |
| # ---------------- UI: TRANSLATION ---------------- | |
| def build_translation_page(model_name): | |
| st.title(f"🌍 {model_name}: Tradução") | |
| if "nllb" in model_name: | |
| lang_map = {"Português": "por_Latn", "Kabuverdianu": "kea_Latn"} | |
| elif "m2m100" in model_name: | |
| lang_map = {"Português": "__pt__", "Kabuverdianu": "__en__"} | |
| else: | |
| lang_map = {"Português": "pt_XX", "Kabuverdianu": "en_XX"} | |
| col1, col2 = st.columns(2) | |
| with col1: src_label = st.selectbox("Língua de Origem", list(lang_map.keys())) | |
| with col2: tgt_label = st.selectbox("Língua de Destino", list(lang_map.keys())) | |
| text = st.text_area("Texto de entrada", "Katxór sta trás di pórta.", height=100) | |
| if st.button("Traduzir"): | |
| if not text.strip(): | |
| st.warning("Introduza texto!") | |
| return | |
| with st.spinner("A traduzir..."): | |
| try: | |
| result = instantiate_translation_model(model_name, text, lang_map[src_label], lang_map[tgt_label]) | |
| st.success("Resultado:") | |
| st.write(result) | |
| except Exception as e: | |
| st.error(f"Erro: {e}") | |
| # ---------------- UI: DECODER ---------------- | |
| def build_decoder_page(model_name, is_peft=False): | |
| st.title(f"✍️ {model_name}: Geração de Texto") | |
| max_length = st.sidebar.slider("Máximo de Tokens", 10, 200, 50) | |
| num_seq = st.sidebar.number_input("Sequências", 1, 5, 1) | |
| text = st.text_area("Prompt", "Katxór sta trás di pórta.") | |
| if st.button("Gerar"): | |
| spinner_text = "A processar... (Modelos grandes podem demorar muito a carregar da primeira vez)" if is_peft else "A processar..." | |
| with st.spinner(spinner_text): | |
| try: | |
| if is_peft: | |
| results = instantiate_peft_decoder(model_name, max_length, num_seq, text) | |
| else: | |
| # Replaced undefined instantiate_gpt2 framework with native fallback pipeline | |
| pipe = load_pipeline("text-generation", f"Iscte-Sintra/{model_name}") | |
| res_raw = pipe(text, max_new_tokens=max_length, num_return_sequences=num_seq, do_sample=True) | |
| results = [{"generated_text": r["generated_text"]} for r in res_raw] | |
| for res in results: | |
| st.info(res["generated_text"]) | |
| except Exception as e: | |
| st.error(f"Erro: {e}") | |
| # ---------------- UI: ENCODER ---------------- | |
| def build_encoder_page(model_name): | |
| st.title(f"🔍 {model_name}: Fill-Mask") | |
| top_k = st.sidebar.slider("Top K sugestões", 1, 5, 3) | |
| mask_token = "<mask>" if "RoBERTa" in model_name else "[MASK]" | |
| st.write(f"Use o token **{mask_token}** para a palavra em falta.") | |
| input_text = st.text_input("Frase", f"Katxór sta trás di {mask_token}.") | |
| if st.button("Prever"): | |
| try: | |
| results = instantiate_encoder(model_name, top_k, input_text) | |
| for res in results: | |
| st.write(f"✅ **{res['token_str']}** ({res['score']:.2%})") | |
| except Exception: | |
| st.error(f"Certifique-se que usou o token {mask_token}") | |
| # ---------------- MAIN ---------------- | |
| model_dict = { | |
| "RoBERTa-Kriolu": "Encoder", | |
| "GPT2_v1.18": "Decoder", | |
| "LLM-kea-v1.0": "Decoder", | |
| "Modelo-Traducao-kea-ptpt-v1.0": "Encoder-Decoder", | |
| "LLAMA-1B-Kea-CPT": "Decoder", | |
| "nllb-v1.0": "Encoder-Decoder", | |
| "m2m100-v1.0": "Encoder-Decoder", | |
| "mbart-v2.0": "Encoder-Decoder", | |
| "m2m100-v2.0": "Encoder-Decoder", | |
| "mbart-v2.1": "Encoder-Decoder", | |
| "nllb-v2.0": "Encoder-Decoder", | |
| "m2m100-v2.1": "Encoder-Decoder", | |
| "mbart-v2.3": "Encoder-Decoder", | |
| "nllb-kea-v2.1": "Encoder-Decoder", | |
| "portuguese-experiment": "Chatbot", | |
| "opus-v1.1": "Encoder-Decoder", | |
| "mbart-v2.5": "Encoder-Decoder", | |
| "nllb-kea-v2.3": "Encoder-Decoder", | |
| "llama-70B-kea-B": "PEFT-Chatbot", # Updated architecture flag | |
| "llama-3-70b-kriolo-cpt": "PEFT-Chatbot", # Updated architecture flag | |
| "llama-3B-kriolu-cpt": "Encoder-Decoder", | |
| "nllb-v2.5": "Encoder-Decoder", | |
| "mbart-v2.6": "Encoder-Decoder", | |
| "m2m100-v2.5" : "Encoder-Decoder", | |
| "opus-v1.2": "Encoder-Decoder", | |
| "kea_chatbot": "Chatbot" | |
| } | |
| selected_model = st.sidebar.selectbox("Escolha o Modelo", list(model_dict.keys())) | |
| arch = model_dict[selected_model] | |
| if arch == "Chatbot": | |
| build_chatbot_page(selected_model, is_peft=False) | |
| elif arch == "PEFT-Chatbot": | |
| build_chatbot_page(selected_model, is_peft=True) | |
| elif arch == "Encoder": | |
| build_encoder_page(selected_model) | |
| elif arch == "Encoder-Decoder": | |
| build_translation_page(selected_model) | |
| elif arch == "PEFT-Decoder": | |
| build_decoder_page(selected_model, is_peft=True) | |
| else: | |
| build_decoder_page(selected_model, is_peft=False) |