GitBot / app.py
acecalisto3's picture
Update app.py
2ec88b8 verified
raw
history blame
3.91 kB
import gradio as gr
import requests
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from sentence_transformers import SentenceTransformer, util
# Initialize models and tokenizers
model_name = "enricoros/big-agi"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Constants
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):
# Generate a response using the loaded model
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
response = generator(issue_text, max_length=512, num_return_sequences=1)[0]['generated_text']
return response
def find_related_issues(issue_text: str, issues: list):
issue_embedding = similarity_model.encode(issue_text)
related_issues = []
for issue in issues:
title_embedding = similarity_model.encode(issue['title'])
similarity = util.cos_sim(issue_embedding, title_embedding)[0][0]
related_issues.append((issue, similarity.item()))
related_issues.sort(key=lambda x: x[1], reverse=True)
return [issue for issue, _ in related_issues[:MAX_RELATED_ISSUES]]
def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str):
headers = {'Authorization': f'token {github_api_token}'}
url = f"{GITHUB_API_BASE_URL}/{github_username}/{github_repository}/issues"
response = requests.get(url, headers=headers)
if response.status_code == 200:
return response.json()
else:
raise Exception(f"Failed to fetch issues: {response.text}")
def respond(
command,
history,
system_message,
max_tokens,
temperature,
top_p,
github_api_token,
github_username,
github_repository,
selected_model,
severity,
programming_language,
):
# Processing the command and generating a response
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
response = generator(f"{system_message}\n{command}\n{history}", max_length=max_tokens, num_return_sequences=1)[0]['generated_text']
return response
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.", label="System message")
model_dropdown = gr.Dropdown(choices=[DEFAULT_MODEL, "enricoros/big-agi"], label="Select Model for Issue Resolution", value=DEFAULT_MODEL)
severity_dropdown = gr.Dropdown(choices=["Critical", "Major", "Minor", "Trivial"], label="Severity")
programming_language_textbox = gr.Textbox(label="Programming Language")
command_dropdown = gr.Dropdown(choices=["/github", "/help", "/generate_code"], label="Select Command")
chatbot = gr.Interface(
fn=respond,
inputs=[
command_dropdown,
system_message,
gr.Slider(minimum=1, maximum=8192, value=2048, label="Max new tokens"),
gr.Slider(minimum=0.1, maximum=4.0, value=0.71, label="Temperature"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p (nucleus sampling)"),
github_api_token,
github_username,
github_repository,
model_dropdown,
severity_dropdown,
programming_language_textbox
],
outputs="text"
)
if __name__ == "__main__":
demo.launch(share=True, server_name="0.0.0.0", server_port=7860)