Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from transformers import (
     AutoTokenizer,
-    AutoModelForSeq2SeqLM,
+    AutoModelForCausalLM,  # Use AutoModelForCausalLM instead of AutoModelForSeq2SeqLM
     TextIteratorStreamer
 )
 from huggingface_hub import login
@@ -18,7 +18,7 @@ if not HF_TOKEN:
     raise ValueError("Missing Hugging Face Token. Please set the HF_TOKEN environment variable.")
 
 # ✅ Only PT-T5 Model
-MODEL_NAME = "amiguel/Meta-Llama-3.1-8B-Instruct-lei-geral-trabalho"# "amiguel/mistral-angolan-laborlaw" #"amiguel/mistral-angolan-laborlaw-ptt5"
+MODEL_NAME = "amiguel/Meta-Llama-3.1-8B-Instruct-lei-geral-trabalho"
 
 # UI Setup
 st.set_page_config(page_title="Assistente LGT | Angola", page_icon="🚀", layout="centered")
@@ -54,7 +54,7 @@ def load_model():
     try:
         login(token=HF_TOKEN)
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN, use_fast=False)
-        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32).to("cuda" if torch.cuda.is_available() else "cpu")
+        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32).to("cuda" if torch.cuda.is_available() else "cpu")
         return model, tokenizer
     except Exception as e:
         st.error(f"🤖 Erro ao carregar o modelo: {str(e)}")
@@ -63,10 +63,8 @@ def load_model():
 # Streaming response generation
 def generate_response(prompt, context, model, tokenizer):
     full_prompt = f"Contexto:\n{context}\n\nPergunta: {prompt}\nResposta:"
-
     inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True).to(model.device)
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
     generation_kwargs = {
         "input_ids": inputs["input_ids"],
         "attention_mask": inputs["attention_mask"],
@@ -78,7 +76,6 @@ def generate_response(prompt, context, model, tokenizer):
         "use_cache": True,
         "streamer": streamer
     }
-
     Thread(target=model.generate, kwargs=generation_kwargs).start()
     return streamer
 
@@ -100,7 +97,7 @@ if prompt := st.chat_input("Faca uma pergunta sobre a LGT..."):
 
     # Load model if not loaded
     if "model" not in st.session_state:
-        with st.spinner("🔄 A carregar o modelo PT-T5..."):
+        with st.spinner("🔄 A carregar o modelo ..."):
             model, tokenizer = load_model()
             if not model:
                 st.stop()
@@ -119,25 +116,20 @@ if prompt := st.chat_input("Faca uma pergunta sobre a LGT..."):
         try:
             start_time = time.time()
             streamer = generate_response(prompt, context, model, tokenizer)
-
             for chunk in streamer:
                 full_response += chunk.strip() + " "
                 response_box.markdown(full_response + "▌", unsafe_allow_html=True)
-
             end_time = time.time()
             input_tokens = len(tokenizer(prompt)["input_ids"])
             output_tokens = len(tokenizer(full_response)["input_ids"])
             speed = output_tokens / (end_time - start_time)
-            cost_usd = ((input_tokens / 1e6) * 5) + ((output_tokens / 1e6) * 15)
+            cost_usd = ((input_tokens / 1e6) * 0.0001) + ((output_tokens / 1e6) * 0.0001)
             cost_aoa = cost_usd * 1160
-
             st.caption(
                 f"🔑 Tokens: {input_tokens} → {output_tokens} | 🕒 Velocidade: {speed:.1f}t/s | "
                 f"💰 USD: ${cost_usd:.4f} | 🇦🇴 AOA: {cost_aoa:.2f}"
             )
-
             response_box.markdown(full_response.strip())
             st.session_state.messages.append({"role": "assistant", "content": full_response.strip()})
-
         except Exception as e:
-            st.error(f"⚡ Erro ao gerar resposta: {str(e)}")
+            st.error(f"⚡ Erro ao gerar resposta: {str(e)}")