fhariya/Legal_document_question_and_answer

1import gradio as gr2import torch3from transformers import Autotokenizer, LlamaForCasualLM, Llamaconfig4def answer_query(query):5    try:6        input_prompt = f"Answer the following question concisely: {query}"7        inputs = tokenizer(input_prompt, return_tensors="pt").to(model.device)8 9        # Generate a response10        with torch.no_grad():11            output = model.generate(12                **inputs,13                max_length=150,  # Limit max length to reduce response time14                num_return_sequences=1,15                no_repeat_ngram_size=2,16                early_stopping=True,17                temperature=0.5,18                top_p=0.8,19                num_beams=120            )21 22        # Decode the generated response23        response = tokenizer.decode(output[0], skip_special_tokens=True)24 25 26        response = response.replace(input_prompt, "").strip()27 28        if query in response:29            response = response.replace(query, "").strip()30 31        # Check if response ends abruptly32        if response.endswith(('.', '!', '?')) is False:33            response += " (The answer may continue, please ask for more details if needed.)"34 35        return response36 37    except RuntimeError as e:38        # Catch CUDA out-of-memory errors39        if "out of memory" in str(e):40            torch.cuda.empty_cache()41            return "Error: Out of memory. Try simplifying the question or reducing the context length."42        else:43            return f"Error: {str(e)}"44 45# Define the Gradio interface46iface = gr.Interface(47    fn=answer_query,48    inputs="text",49    outputs="text",50    title="Legal Document Question Answering",51    description="Hey Lawyer! Ask questions about the legal documents.",52)53 54# Using vertical block layout55with gr.Blocks() as demo:56    gr.Markdown("# Legal Document Question Answering")57    gr.Markdown("### Hey Lawyer! Ask questions about the legal documents below.")58 59    with gr.Column():60        query_input = gr.Textbox(label="Your Question:")61        submit_btn = gr.Button("Submit", variant="primary")62        clear_btn = gr.Button("Clear", variant="secondary")63        response_output = gr.Textbox(label="Answer:", interactive=False)64 65    submit_btn.click(answer_query, inputs=query_input, outputs=response_output)66    clear_btn.click(lambda: ("", ""), inputs=None, outputs=[query_input, response_output])67 68iface.launch(inline=False)69