Update app.py

app.py

@@ -1,17 +1,14 @@
 import gradio as gr
-from transformers import pipeline
-from sentence_transformers import SentenceTransformer, util
-import os
 import requests
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from sentence_transformers import SentenceTransformer, util
 
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-
+# Initialize models and tokenizers
 model_name = "enricoros/big-agi"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-
-# Constants for enhanced organization
+# Constants
 GITHUB_API_BASE_URL = "https://api.github.com/repos"
 DEFAULT_MODEL = "microsoft/CodeBERT-base"
 MAX_RELATED_ISSUES = 3
@@ -19,58 +16,32 @@ 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()
+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']
 
-    # 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()
-
-    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
+def find_related_issues(issue_text: str, issues: list):
     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]]
+    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()))
 
-    return related_issues
+    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) -> list:
-    # Fetch the issues from the GitHub API
+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)
-
-    # Parse the JSON response
-    issues = response.json()
-
-    return issues
+    if response.status_code == 200:
+        return response.json()
+    else:
+        raise Exception(f"Failed to fetch issues: {response.text}")
 
 def respond(
     command,
@@ -85,29 +56,11 @@ def respond(
     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,
-    }
+):
+    # 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():
@@ -115,66 +68,29 @@ 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. 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,
-    )
-
+    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")
 
-    command_dropdown = gr.Dropdown(
-        choices=[
-            "/github",
-            "/help",
-            "/generate_code",
-            "/explain_concept",
-            "/write_documentation",
-            "/translate_code",
-        ],
-        label="Select Command",
-    )
-
-    chatbot = MyChatbot(
-        respond,
-        additional_inputs=[
+    chatbot = gr.Interface(
+        fn=respond,
+        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)",
-            ),
+            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,
+            programming_language_textbox
         ],
+        outputs="text"
     )
 
 if __name__ == "__main__":
-    demo.queue().launch(
-        share=True,
-        server_name="0.0.0.0",
-        server_port=7860
-    )
+    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)