Spaces:

albhu
/

legalgeek

Running

App Files Files Community

legalgeek / app.py

albhu

Update app.py

aafe73b verified about 1 month ago

raw history blame contribute delete

No virus

3.41 kB

	import streamlit as st
	from pdfminer.high_level import extract_text
	import docx
	from transformers import AutoTokenizer, AutoModelForCausalLM

	st.title("JAI")

	# Előző beszélgetések tárolása session state-ben
	if "previous_conversations" not in st.session_state:
	st.session_state.previous_conversations = []

	# Korábbi beszélgetések törlése
	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 = AutoTokenizer.from_pretrained("pile-of-law/legalbert-large-1.7M-2")
	model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-3b", trust_remote_code=True)

	# Dokumentumfeldolgozó függvény
	def process_document(document_file):
	document_text = ""
	if document_file.type == "application/pdf":
	document_text = extract_text(document_file)
	elif document_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
	docx_file = docx.Document(document_file)
	for paragraph in docx_file.paragraphs:
	text = paragraph.text
	if text:
	document_text += text.strip() + "\n\n"
	elif document_file.type == "application/msword":
	doc_file = docx.Document(document_file)
	for paragraph in doc_file.paragraphs:
	text = paragraph.text
	if text:
	document_text += text.strip() + "\n\n"
	else:
	st.error("A fájltípus nem támogatott. Kérlek válassz ki egy PDF, DOCX vagy DOC fájlt!")
	return document_text

	# Válaszgeneráló függvény
	def generate_response(input_text):
	# Bemeneti szöveg tisztítása
	cleaned_input = input_text.strip()

	# Tokenizálás
	inputs = tokenizer(cleaned_input, return_tensors="pt", max_length=1024, truncation=True)

	# Modell használata a válasz generálásához
	outputs = model(**inputs)

	# Válasz visszaadása
	response = tokenizer.decode(outputs.logits.argmax(dim=1)[0])
	return response

	# Dokumentumfeltöltés feldolgozása
	document_file = st.file_uploader("Húzd ide a dokumentumot vagy kattints a feltöltéshez", type=["pdf", "docx", "doc"])

	if document_file is not None:
	document_text = process_document(document_file)

	# Előző beszélgetésekhez csatolható kontextus
	context = st.text_area("Korábbi Beszélgetéshez Tartozó Kontextus", "")

	# Kérdés mező hozzáadása
	question = st.text_input("Kérdés a Dokumentumból", "")

	# Ha van kérdés, azt is hozzáadjuk a bemenethez
	if question:
	input_text = f"{context} {document_text} {question}" if context else f"{document_text} {question}"
	else:
	input_text = f"{context} {document_text}" if context else document_text

	# Válasz generálása csak akkor, ha a felhasználó elküldi a promptot
	if input_text.strip() != "":
	response = generate_response(input_text)
	st.subheader("Generált Válasz:")
	st.write(response)

	# Aktuális beszélgetés hozzáadása az előző beszélgetésekhez
	st.session_state.previous_conversations.append({"input_text": input_text, "response": response})

	# Input text megjelenítése
	st.subheader("Input Text:")
	st.write(input_text)

	# Response megjelenítése
	st.subheader("Response:")
	st.write(response)