Update app.py

app.py

@@ -1,7 +1,9 @@
 import streamlit as st
+import pdfplumber
+import docx
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-st.title("Jogi Dokumentumok AI Válaszgenerálás")
+st.title("JAI")
 
 # Előző beszélgetések tárolása session state-ben
 if "previous_conversations" not in st.session_state:
@@ -12,22 +14,19 @@ if st.checkbox("Korábbi Beszélgetések Törlése"):
     st.session_state.previous_conversations = []
 
 # Tokenizáló modell és válaszgeneráló modell választása
-tokenizer_name = "pile-of-law/legalbert-large-1.7M-2"
-model_name = "apple/OpenELM-3B-Instruct"
-
-tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("AdaptLLM/law-chat")
+model = AutoModelForCausalLM.from_pretrained("AdaptLLM/law-chat", trust_remote_code=True)
 
 # Definiáljuk a válaszgeneráló függvényt
 def generate_response(input_text, tokenizer, model):
     # Tokenizálás
-    input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=tokenizer.model_max_length, truncation=True)
+    inputs = tokenizer(input_text, return_tensors="pt", max_length=1024)
 
     # Modell használata a válasz generálásához
-    output = model.generate(input_ids, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+    outputs = model(**inputs)
 
-    # Decoding és válasz visszaadása
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    # Válasz visszaadása
+    response = tokenizer.decode(outputs.logits.argmax(dim=1)[0])
     return response
 
 # Dokumentum feltöltése drag and drop segítségével
@@ -37,18 +36,21 @@ if document_file is not None:
     document_text = ""
 
     if document_file.type == "application/pdf":
+        # PDF fájl feldolgozása
         with pdfplumber.open(document_file) as pdf:
             for page in pdf.pages:
                 document_text += page.extract_text()
                 # Ha bekezdésekre bontjuk, hozzáadhatunk például egy sortörést a bekezdések közé:
                 document_text += "\n\n"
     elif document_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+        # DOCX fájl feldolgozása
         docx_file = docx.Document(document_file)
         for paragraph in docx_file.paragraphs:
             document_text += paragraph.text
             # Hasonlóan, itt is adhatunk hozzá sortörést a bekezdések közé:
             document_text += "\n\n"
     elif document_file.type == "application/msword":
+        # DOC fájl feldolgozása
         doc_file = docx.Document(document_file)
         for paragraph in doc_file.paragraphs:
             document_text += paragraph.text