Plong95's picture
Upload 3 files
fa62358
import gradio as gr
import git
import os
from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
def clone_repo(repo_url):
local_path = "repo_clone"
git.Repo.clone_from(repo_url, local_path)
return local_path
def process_repo(repo_url, option):
if option == "Pre-trained":
qa_pipeline = pipeline('question-answering')
else:
model_path = "./model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForQuestionAnswering.from_pretrained(model_path)
qa_pipeline = pipeline('question-answering', model=model, tokenizer=tokenizer)
repo_path = clone_repo(repo_url)
result = {}
for root, dirs, files in os.walk(repo_path):
for file in files:
file_path = os.path.join(root, file)
with open(file_path, 'r', encoding="utf-8") as f:
text = f.read()
summary = text[:50] + "..." if len(text) > 50 else text
keywords = qa_pipeline(summary)['answer']
result[file_path] = {"summary": summary, "text": text, "keywords": keywords}
return result
def qa_chatbot(repo_dict, question):
all_text = ""
for file in repo_dict.values():
all_text += file['summary'] + " " + file['text'] + " "
answer = qa_pipeline({'context': all_text, 'question': question})['answer']
return answer
input_repo = gr.inputs.Textbox(label="Enter Git repository URL")
output_processed_repo = gr.outputs.Textbox(label="Processed Git repository")
output_qa_chatbot = gr.outputs.Textbox(label="Answer")
model_options = ["Pre-trained", "Fine-tuned"]
input_option = gr.inputs.Dropdown(choices=model_options, label="Choose a model option")
process_repo_interface = gr.Interface(fn=process_repo, inputs=[input_repo, input_option], outputs=output_processed_repo,
title="Process Git Repository")
qa_chatbot_interface = gr.Interface(fn=qa_chatbot, inputs={"repo_dict": gr.inputs.Dictionary(
key_type=gr.inputs.Textbox(label="File path"),
value_type=gr.inputs.Dictionary(key_type=gr.inputs.Textbox(label="File content"),
value_type=gr.inputs.Textbox(label="Keywords"))),
"question": gr.inputs.Textbox(label="Question")},
outputs=output_qa_chatbot,
title="QA Chatbot")
process_repo_interface.launch()
qa_chatbot_interface.launch()