Create app.py

app.py

@@ -0,0 +1,88 @@
+import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForQuestionAnswering
+
+def carregar_modelos():
+    # Caminhos no Hugging Face Hub
+    caminho_gpt = "mateuscasare/GPTNeo-Vade-Mecum-2023"
+    caminho_bert = "mateuscasare/BERT-Vade-Mecum-2023"
+    caminho_bert_juridico = "mateuscasare/modelo_juridico_pt"
+
+    dispositivo = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Carregar modelos do Hugging Face Hub
+    tokenizer_gpt = AutoTokenizer.from_pretrained(caminho_gpt)
+    model_gpt = AutoModelForCausalLM.from_pretrained(caminho_gpt).to(dispositivo)
+
+    tokenizer_bert = AutoTokenizer.from_pretrained(caminho_bert)
+    model_bert = AutoModelForQuestionAnswering.from_pretrained(caminho_bert).to(dispositivo)
+
+    tokenizer_bert_juridico = AutoTokenizer.from_pretrained(caminho_bert_juridico)
+    model_bert_juridico = AutoModelForQuestionAnswering.from_pretrained(caminho_bert_juridico).to(dispositivo)
+
+    return tokenizer_gpt, model_gpt, tokenizer_bert, model_bert, tokenizer_bert_juridico, model_bert_juridico, dispositivo
+
+def main():
+    st.set_page_config(page_title="Assistente Jurídico", layout="wide")
+    st.title("⚖️ Assistente Jurídico")
+
+    # Carregar modelos
+    tokenizer_gpt, model_gpt, tokenizer_bert, model_bert, tokenizer_bert_juridico, model_bert_juridico, dispositivo = carregar_modelos()
+
+    # Interface principal
+    tab1, tab2, tab3 = st.tabs(["Geração de Texto", "Perguntas e Respostas", "Perguntas Jurídicas Específicas"])
+
+    with tab1:
+        st.subheader("📝 Gerador de Textos Jurídicos")
+        prompt = st.text_area("Descreva sua necessidade:", height=150, placeholder="Ex: Preciso de uma petição de habeas corpus...")
+
+        col1, col2 = st.columns(2)
+        with col1:
+            max_length = st.slider("Tamanho da resposta", 100, 2000, 500)
+            temperatura = st.slider("Criatividade", 0.0, 1.0, 0.7)
+
+        if st.button("Gerar", type="primary"):
+            with st.spinner("Gerando texto..."):
+                inputs = tokenizer_gpt(prompt, return_tensors="pt").to(dispositivo)
+                outputs = model_gpt.generate(**inputs, max_length=max_length, temperature=temperatura, do_sample=True)
+                texto_gerado = tokenizer_gpt.decode(outputs[0], skip_special_tokens=True)
+                st.text_area("Texto gerado:", texto_gerado, height=400)
+
+    with tab2:
+        st.subheader("❓ Consulta Jurídica (Modelo Genérico)")
+        contexto = st.text_area("Contexto:", height=150, placeholder="Insira o contexto jurídico aqui...")
+        pergunta = st.text_input("Sua pergunta:", placeholder="Ex: Qual o prazo para...?")
+
+        if st.button("Responder (Modelo Genérico)", type="primary"):
+            with st.spinner("Analisando..."):
+                inputs = tokenizer_bert(pergunta, contexto, return_tensors="pt").to(dispositivo)
+                outputs = model_bert(**inputs)
+                answer_start = torch.argmax(outputs.start_logits)
+                answer_end = torch.argmax(outputs.end_logits) + 1
+                answer = tokenizer_bert.decode(inputs["input_ids"][0][answer_start:answer_end])
+                st.success(f"Resposta: {answer}")
+
+    with tab3:
+        st.subheader("⚖️ Consulta Jurídica Específica")
+        contexto_juridico = st.text_area("Contexto Jurídico:", height=150, placeholder="Insira o contexto jurídico específico aqui...")
+        pergunta_juridica = st.text_input("Sua pergunta jurídica:", placeholder="Ex: Qual o prazo para recorrer?")
+
+        if st.button("Responder (Modelo Jurídico)", type="primary"):
+            with st.spinner("Analisando..."):
+                inputs = tokenizer_bert_juridico(pergunta_juridica, contexto_juridico, return_tensors="pt").to(dispositivo)
+                outputs = model_bert_juridico(**inputs)
+                answer_start = torch.argmax(outputs.start_logits)
+                answer_end = torch.argmax(outputs.end_logits) + 1
+                answer = tokenizer_bert_juridico.decode(inputs["input_ids"][0][answer_start:answer_end])
+                st.success(f"Resposta Jurídica: {answer}")
+
+    with st.expander("ℹ️ Sobre os Modelos"):
+        st.markdown("""
+        Este assistente utiliza três modelos de IA:
+        - **GPTNeo Vade-Mecum**: Para geração de textos jurídicos.
+        - **BERT Vade-Mecum**: Para responder perguntas gerais sobre direito.
+        - **Modelo Jurídico PT**: Treinado especificamente com textos jurídicos brasileiros para maior precisão.
+        """)
+
+if __name__ == "__main__":
+    main()