Spaces:
Runtime error
Runtime error
File size: 4,717 Bytes
6e424ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
import streamlit as st
import os
import re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from PyPDF2 import PdfReader
from peft import get_peft_model, LoraConfig, TaskType
# β
Force CPU execution
device = torch.device("cpu")
# πΉ Load IBM Granite Model (CPU-Compatible)
MODEL_NAME = "ibm-granite/granite-3.1-2b-instruct"
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map="cpu", # Force CPU execution
torch_dtype=torch.float32 # Use float32 since Hugging Face runs on CPU
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# πΉ Apply LoRA Fine-Tuning Configuration
lora_config = LoraConfig(
r=8,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.1,
bias="none",
task_type=TaskType.CAUSAL_LM
)
model = get_peft_model(model, lora_config)
model.eval()
# π Function to Read & Extract Text from PDFs
def read_files(file):
file_context = ""
reader = PdfReader(file)
for page in reader.pages:
text = page.extract_text()
if text:
file_context += text + "\n"
return file_context.strip()
# π Function to Format AI Prompts
def format_prompt(system_msg, user_msg, file_context=""):
if file_context:
system_msg += f" The user has provided a contract document. Use its context to generate insights, but do not repeat or summarize the document itself."
return [
{"role": "system", "content": system_msg},
{"role": "user", "content": user_msg}
]
# π Function to Generate AI Responses
def generate_response(input_text, max_tokens=1000, top_p=0.9, temperature=0.7):
model_inputs = tokenizer([input_text], return_tensors="pt").to(device)
with torch.no_grad():
output = model.generate(
**model_inputs,
max_new_tokens=max_tokens,
do_sample=True,
top_p=top_p,
temperature=temperature,
num_return_sequences=1,
pad_token_id=tokenizer.eos_token_id
)
return tokenizer.decode(output[0], skip_special_tokens=True)
# π Function to Clean AI Output
def post_process(text):
cleaned = re.sub(r'ζ₯+', '', text) # Remove unwanted symbols
lines = cleaned.splitlines()
unique_lines = list(dict.fromkeys([line.strip() for line in lines if line.strip()]))
return "\n".join(unique_lines)
# π Function to Handle RAG with IBM Granite & Streamlit
def granite_simple(prompt, file):
file_context = read_files(file) if file else ""
system_message = "You are IBM Granite, a legal AI assistant specializing in contract analysis."
messages = format_prompt(system_message, prompt, file_context)
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
response = generate_response(input_text)
return post_process(response)
# πΉ Streamlit UI
def main():
st.set_page_config(page_title="Contract Analysis AI", page_icon="π")
st.title("π AI-Powered Contract Analysis Tool")
st.write("Upload a contract document (PDF) for a detailed AI-driven legal and technical analysis.")
# πΉ Sidebar Settings
with st.sidebar:
st.header("βοΈ Settings")
max_tokens = st.slider("Max Tokens", 50, 1000, 250, 50)
top_p = st.slider("Top P (sampling)", 0.1, 1.0, 0.9, 0.1)
temperature = st.slider("Temperature (creativity)", 0.1, 1.0, 0.7, 0.1)
# πΉ File Upload Section
uploaded_file = st.file_uploader("π Upload a contract document (PDF)", type="pdf")
# β
Ensure file upload message is displayed
if uploaded_file is not None:
st.session_state["uploaded_file"] = uploaded_file # Persist file in session state
st.success("β
File uploaded successfully!")
st.write("Click the button below to analyze the contract.")
# Force button to always render
st.markdown('<style>div.stButton > button {display: block; width: 100%;}</style>', unsafe_allow_html=True)
if st.button("π Analyze Document"):
with st.spinner("Analyzing contract document... β³"):
final_answer = granite_simple(
"Perform a detailed technical analysis of the attached contract document, highlighting potential risks, legal pitfalls, compliance issues, and areas where contractual terms may lead to future disputes or operational challenges.",
uploaded_file
)
# πΉ Display Analysis Result
st.subheader("π Analysis Result")
st.write(final_answer)
# π₯ Run Streamlit App
if __name__ == '__main__':
main()
|