GitBot / app.py
acecalisto3's picture
Update app.py
a89645e verified
raw
history blame
5.53 kB
import gradio as gr
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util
import os
import requests
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
model_name = "enricoros/big-agi"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Constants for enhanced organization
GITHUB_API_BASE_URL = "https://api.github.com/repos"
DEFAULT_MODEL = "microsoft/CodeBERT-base"
MAX_RELATED_ISSUES = 3
# Load a pre-trained model for sentence similarity
similarity_model = SentenceTransformer('all-mpnet-base-v2')
def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None) -> str:
# Initialize the model
model = pipeline("text-generation", model=model_name)
# Generate a response
response = model(
f"{system_message}\n{issue_text}\nAssistant: ",
max_length=max_tokens,
do_sample=True,
temperature=temperature,
top_k=top_p,
)
# Extract the assistant's response
assistant_response = response[0]['generated_text'].strip()
# Analyze the response
if "Severity" in assistant_response:
severity = assistant_response.split(":")[1].strip()
if "Programming Language" in assistant_response:
programming_language = assistant_response.split(":")[1].strip()
return {
'assistant_response': assistant_response,
'severity': severity,
'programming_language': programming_language,
}
def find_related_issues(issue_text: str, issues: list) -> list:
# Calculate the similarity between the issue and other issues
issue_embedding = similarity_model.encode(issue_text)
similarities = [util.cos_sim(issue_embedding, similarity_model.encode(issue['title'])) for issue in issues]
# Sort the issues by similarity
sorted_issues = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
# Select the top related issues
related_issues = [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]
return related_issues
def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str) -> list:
# Fetch the issues from the GitHub API
headers = {'Authorization': f'token {github_api_token}'}
url = f"{GITHUB_API_BASE_URL}/{github_username}/{github_repository}/issues"
response = requests.get(url, headers=headers)
# Parse the JSON response
issues = response.json()
return issues
def respond(
command,
history,
system_message,
max_tokens,
temperature,
top_p,
github_api_token,
github_username,
github_repository,
selected_model,
severity,
programming_language,
*args,
**kwargs,
) -> dict:
# Initialize the model
model = pipeline("text-generation", model="enricoros/big-agi")
# Generate a response
response = model(
f"{system_message}\n{command}\n{history}\n{github_username}/{github_repository}\n{severity}\n{programming_language}\nAssistant: ",
max_length=max_tokens,
do_sample=True,
temperature=temperature,
top_k=top_p,
)
# Extract the assistant's response
assistant_response = response[0]['generated_text'].strip()
return {
'assistant_response': assistant_response,
'severity': severity,
'programming_language': programming_language,
}
)
with gr.Blocks() as demo:
with gr.Row():
github_api_token = gr.Textbox(label="GitHub API Token", type="password")
github_username = gr.Textbox(label="GitHub Username")
github_repository = gr.Textbox(label="GitHub Repository")
system_message = gr.Textbox(
value="You are GitBot, the Github project guardian angel. You resolve issues and propose implementation of feature requests",
label="System message",
)
model_dropdown = gr.Dropdown(
choices=[
"microsoft/CodeBERT-base",
"Salesforce/codegen-350M-mono",
],
label="Select Model for Issue Resolution",
value=DEFAULT_MODEL,
)
severity_dropdown = gr.Dropdown(
choices=["Critical", "Major", "Minor", "Trivial"],
label="Severity",
value=None,
)
programming_language_textbox = gr.Textbox(label="Programming Language")
command_dropdown = gr.Dropdown(
choices=[
"/github",
"/help",
"/generate_code",
"/explain_concept",
"/write_documentation",
"/translate_code",
],
label="Select Command",
)
chatbot = MyChatbot(
respond,
additional_inputs=[
command_dropdown,
system_message,
gr.Slider(minimum=1, maximum=8192, value=2048, step=1, label="Max new tokens"),
gr.Slider(minimum=0.1, maximum=4.0, value=0.71, step=0.1, label="Temperature"),
gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.95,
step=0.01,
label="Top-p (nucleus sampling)",
),
github_api_token,
github_username,
github_repository,
model_dropdown,
severity_dropdown,
programming_language_textbox,
],
)
if __name__ == "__main__":
demo.queue().launch(
share=True,
server_name="0.0.0.0",
server_port=7860
)