Update app.py

app.py

@@ -1,57 +1,40 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from sentence_transformers import SentenceTransformer
-import faiss
-import numpy as np
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-# Generatív Modell Betöltése
+# Modell betöltése
 tokenizer = AutoTokenizer.from_pretrained("SZTAKI-HLT/hubert-base-cc")
-gpt_model = AutoModelForSeq2SeqLM.from_pretrained("SZTAKI-HLT/hubert-base-cc")
-generate_answer = pipeline("text2text-generation", model=gpt_model, tokenizer=tokenizer)
-
-# FAISS index létrehozása (az esetleges további funkciókhoz)
-dimension = 384  # Paraphrase-MiniLM-L6-v2 modell embedding dimenziója
-embedder = SentenceTransformer("sentence-transformers/paraphrase-MiniLM-L6-v2")
-index = faiss.IndexFlatL2(dimension)
-document_chunks = []
-chunk_embeddings = []
+model = AutoModelForSeq2SeqLM.from_pretrained("SZTAKI-HLT/hubert-base-cc")
 
 # Streamlit alkalmazás
-st.title("AI Jogi Chatbot - Dokumentumalapú Válaszadás ChatGPT-stílusban")
+st.title("AI Jogi Chatbot Magyar Nyelven")
 
-# Dokumentum feltöltése és indexelése
-uploaded_file = st.file_uploader("Töltsön fel egy dokumentumot", type=["txt", "pdf"])
+# Dokumentum feltöltése
+uploaded_file = st.file_uploader("Töltsön fel egy dokumentumot", type=["txt"])
 
 if uploaded_file:
     try:
-        # Próbáljuk meg UTF-8 kódolással olvasni
+        # Szöveg beolvasása
         context = uploaded_file.read().decode("utf-8")
     except UnicodeDecodeError:
-        # Ha nem sikerül, próbáljuk meg ISO-8859-1 kódolással
         context = uploaded_file.read().decode("ISO-8859-1")
     
     # Teljes dokumentum megjelenítése
     st.write("### Feltöltött dokumentum tartalma:")
     st.write(context)
-
-    # Szöveg feldarabolása és embedding létrehozása FAISS indexhez (opcionális a későbbi használathoz)
-    paragraphs = context.split("\n\n")
-    for paragraph in paragraphs:
-        if len(paragraph.strip()) > 0:
-            embedding = embedder.encode(paragraph).reshape(1, -1)
-            document_chunks.append(paragraph)
-            chunk_embeddings.append(embedding)
-            index.add(embedding)
     
-    st.write(f"A dokumentum sikeresen indexelve. {len(paragraphs)} bekezdés feldolgozva.")
-
     # Chatelés a dokumentum kontextusában
     st.write("### Chat")
     question = st.text_input("Írja be a kérdését a dokumentumról:")
-
+    
     if question:
-        # Dokumentum-alapú generatív válaszadás
-        input_text = f"Kérdés: {question}\nDokumentum tartalma: {context}"
-        response = generate_answer(input_text, max_length=150)[0]["generated_text"]
-
+        with st.spinner('Válasz generálása...'):
+            # Bemenet összeállítása
+            input_text = f"{context}\n\nKérdés: {question}\nVálasz:"
+            inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
+            
+            # Válasz generálása
+            outputs = model.generate(inputs, max_length=150, num_beams=4, early_stopping=True)
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
         st.write("Válasz:", response)