IngenierIA / app.py
milcio79's picture
Update app.py
575ab3d verified
Raw
History Blame Contribute Delete
3.99 kB
# 2. Import necessary libraries
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
import os
# 3. Load the base model and tokenizer from Hugging Face
model_name = "EleutherAI/gpt-neo-1.3B" # You can replace this with another LLM available for free
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# 4. Function to tokenize, vectorize, and find similar parts of a document
def find_relevant_parts(document_text, prompt, top_n=5):
vectorizer = TfidfVectorizer(stop_words='english')
doc_parts = document_text.split('.') # Split document into parts
vectors = vectorizer.fit_transform(doc_parts)
query_vec = vectorizer.transform([prompt])
similarities = (vectors * query_vec.T).toarray().flatten()
top_indices = similarities.argsort()[-top_n:][::-1] # Top `top_n` similar parts
relevant_parts = ' '.join([doc_parts[i] for i in top_indices])
return relevant_parts
# 5. Function for generating a response from the model
def generate_response(document_path, prompt):
with open(document_path, 'r') as file:
document_text = file.read()
relevant_parts = find_relevant_parts(document_text, prompt)
complete_input = prompt + ' ' + relevant_parts
inputs = tokenizer(complete_input, return_tensors="pt")
outputs = model.generate(**inputs, max_length=150, num_return_sequences=1)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
# 6. Gradio Interface setup for user authentication and document querying
def user_authentication(username, password):
# Mock simple authentication logic (Replace with secure user authentication if needed)
valid_users = {"user1": "pass1", "user2": "pass2"} # Replace with more secure methods
if username in valid_users and valid_users[username] == password:
return True
return False
# 7. Gradio App Structure
with gr.Blocks() as demo:
# Page 1: Login page
with gr.Tab("Login"):
username = gr.Textbox(label="Username")
password = gr.Textbox(label="Password", type="password")
login_button = gr.Button("Login")
login_output = gr.Textbox(label="Login Status", interactive=False)
def login_check(user, pwd):
if user_authentication(user, pwd):
return "Login successful!", gr.update(visible=False), gr.update(visible=True)
else:
return "Invalid credentials", gr.update(visible=True), gr.update(visible=False)
login_button.click(login_check, inputs=[username, password], outputs=[login_output])
# Page 2: Document selection
with gr.Tab("Select Document", visible=False) as doc_tab:
doc_path = gr.File(label="Upload your document")
select_button = gr.Button("Select Document")
doc_status = gr.Textbox(label="Document Status", interactive=False)
def document_selected(file):
if file:
return "Document selected successfully!", gr.update(visible=False), gr.update(visible=True)
return "No document selected."
select_button.click(document_selected, inputs=[doc_path], outputs=[doc_status])
# Page 3: Query input and response
with gr.Tab("Query Document", visible=False) as query_tab:
prompt = gr.Textbox(label="Enter your prompt")
submit_button = gr.Button("Submit Query")
response_output = gr.Textbox(label="Response", interactive=False)
def process_query(file, user_prompt):
if file and user_prompt:
return generate_response(file.name, user_prompt)
return "Please provide a valid document and prompt."
submit_button.click(process_query, inputs=[doc_path, prompt], outputs=[response_output])
# Launch the Gradio app
demo.launch()