Update app.py

app.py

@@ -1,13 +1,27 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Modell betöltése
-tokenizer = AutoTokenizer.from_pretrained("SZTAKI-HLT/hubert-base-cc")
-model = AutoModelForSeq2SeqLM.from_pretrained("SZTAKI-HLT/hubert-base-cc")
+# Fordító modell betöltése (M2M100 fordítás magyar és angol között)
+translator_tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
+translator_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/m2m100_418M")
+translation_pipeline = pipeline("translation", model=translator_model, tokenizer=translator_tokenizer)
+
+# Generatív modell betöltése (angol nyelven működik)
+gpt_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
+gpt_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
+generate_answer = pipeline("text2text-generation", model=gpt_model, tokenizer=gpt_tokenizer)
+
+# Fordítás funkciók
+def translate_to_english(text):
+    translated_text = translation_pipeline(text, src_lang="hu", tgt_lang="en", max_length=512)[0]['translation_text']
+    return translated_text
+
+def translate_to_hungarian(text):
+    translated_text = translation_pipeline(text, src_lang="en", tgt_lang="hu", max_length=512)[0]['translation_text']
+    return translated_text
 
 # Streamlit alkalmazás
-st.title("AI Jogi Chatbot Magyar Nyelven")
+st.title("AI Jogi Chatbot - Fordítás Alapú Megközelítéssel")
 
 # Dokumentum feltöltése
 uploaded_file = st.file_uploader("Töltsön fel egy dokumentumot", type=["txt"])
@@ -23,18 +37,24 @@ if uploaded_file:
     st.write("### Feltöltött dokumentum tartalma:")
     st.write(context)
     
+    # Dokumentum fordítása angolra
+    with st.spinner('Fordítás angolra...'):
+        context_english = translate_to_english(context)
+    
     # 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:
         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)
+            # Kérdés fordítása angolra
+            question_english = translate_to_english(question)
+            
+            # Generatív válaszadás angol nyelven
+            input_text = f"{context_english}\n\nQuestion: {question_english}\nAnswer:"
+            response_english = generate_answer(input_text, max_length=150)[0]["generated_text"]
             
-            # 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)
+            # Válasz fordítása magyarra
+            response_hungarian = translate_to_hungarian(response_english)
         
-        st.write("Válasz:", response)
+        st.write("Válasz:", response_hungarian)