Update app.py

app.py

@@ -1,14 +1,10 @@
 import gradio as gr
-import requests
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer, util
+import os
+import requests
 
-# Initialize models and tokenizers
-model_name = "enricoros/big-agi"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-# Constants
+# Constants for enhanced organization
 GITHUB_API_BASE_URL = "https://api.github.com/repos"
 DEFAULT_MODEL = "microsoft/CodeBERT-base"
 MAX_RELATED_ISSUES = 3
@@ -16,51 +12,91 @@ 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']
+# Define models for issue analysis
+model_name = "enricoros/big-agi"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None) -> dict:
+    # Initialize the model
+    model = pipeline("text-generation", model=model_name)
+
+    # Generate a response
+    response = model(
+        f"{issue_text}\nAssistant: ",
+        max_length=512,
+        do_sample=True,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.9,
+    )
+
+    # 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()
 
-    return response
+    if "Programming Language" in assistant_response:
+        programming_language = assistant_response.split(":")[1].strip()
 
-def find_related_issues(issue_text: str, issues: list):
+    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)
-    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()))
+    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)
 
-    related_issues.sort(key=lambda x: x[1], reverse=True)
-    return [issue for issue, _ in related_issues[:MAX_RELATED_ISSUES]]
+    # Select the top related issues
+    related_issues = [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]
 
-def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str):
+    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)
-    if response.status_code == 200:
-        return response.json()
-    else:
-        raise Exception(f"Failed to fetch issues: {response.text}")
+
+    # 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,
-):
-    # 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
+    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}\nSeverity: {severity}\nProgramming Language: {programming_language}\nAssistant: ",
+        max_length=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_k=50,
+        top_p=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():
@@ -68,29 +104,53 @@ with gr.Blocks() as demo:
         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")
+    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-45M-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"], label="Select Command")
 
-    chatbot = gr.Interface(
-        fn=respond,
-        inputs=[
+    command_dropdown = gr.Dropdown(
+        choices=[
+            "/github",
+            "/help",
+            "/generate_code",
+            "/explain_concept",
+            "/write_documentation",
+            "/translate_code",
+        ],
+        label="Select Command",
+    )
+
+    chatbot = gr.Chatbot(
+        respond,
+        additional_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)"),
+            gr.Slider(minimum=1, maximum=8192, value=2048, step=1, label="Max new tokens"),
+            gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+            gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p (nucleus sampling)"),
             github_api_token,
             github_username,
             github_repository,
             model_dropdown,
             severity_dropdown,
-            programming_language_textbox
+            programming_language_textbox,
         ],
-        outputs="text"
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)