Update app.py

app.py

@@ -1,60 +1,27 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from datasets import load_dataset
+from models.codet5 import CodeT5
+from models.other_models import OtherModels
+from repositories.github_api import GitHubAPI
 
-# Define model loading function
-def load_model(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    return tokenizer, model
+codet5_model = CodeT5()
+other_models = OtherModels()
+github_api = GitHubAPI()
 
-# Load selected models
-models = {
-    "bigcode/python-stack-v1-functions-filtered-sc2-subset": "bigcode/python-stack-v1-functions-filtered-sc2-subset",
-    "bigcode/python-stack-v1-functions-filtered-sc2": "bigcode/python-stack-v1-functions-filtered-sc2",
-    "muellerzr/python-stack-v1-functions-filtered-llama-3-8B": "muellerzr/python-stack-v1-functions-filtered-llama-3-8B",
-    "TheBloke/Python-Code-13B-GGUF": "TheBloke/Python-Code-13B-GGUF",
-    "replit/replit-code-v1_5-3b": "replit/replit-code-v1_5-3b",
-    "neulab/codebert-python": "neulab/codebert-python"
-}
+def analyze_repository(repo_url):
+    repo_data = github_api.get_repository(repo_url)
+    optimization_results = codet5_model.analyze(repo_data, github_api)
+    bug_hunting_results = other_models.analyze(repo_data, github_api)
+    return optimization_results, bug_hunting_results
 
-# Load selected datasets
-datasets = {
-    "kye/all-huggingface-python-code": "kye/all-huggingface-python-code",
-    "ajibawa-2023/Python-Code-23k-ShareGPT": "ajibawa-2023/Python-Code-23k-ShareGPT",
-    "suvadityamuk/huggingface-transformers-code-dataset": "suvadityamuk/huggingface-transformers-code-dataset"
-}
-
-# Define the function for code generation
-def generate_code(prompt, model_name, dataset_name, temperature, max_length):
-    tokenizer, model = load_model(models[model_name])
-    
-    # Load dataset (for reference, not directly used)
-    dataset = load_dataset(datasets[dataset_name], split="train")
-
-    # Tokenize input prompt
-    inputs = tokenizer(prompt, return_tensors="pt")
-    
-    # Generate output
-    output_ids = model.generate(**inputs, temperature=temperature, max_length=max_length)
-    generated_code = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    
-    return generated_code
-
-# Create Gradio Interface
 iface = gr.Interface(
-    fn=generate_code,
-    inputs=[
-        gr.Textbox(label="Prompt"),
-        gr.Dropdown(label="Model", choices=list(models.keys())),
-        gr.Dropdown(label="Dataset", choices=list(datasets.keys())),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.5),
-        gr.Slider(label="Max Length", minimum=10, maximum=1000, value=200)
+    fn=analyze_repository,
+    inputs=gr.Textbox(lines=1, placeholder="Enter GitHub Repository URL (e.g., https://github.com/owner/repo)"),
+    outputs=[
+        gr.Textbox(lines=10, label="Optimization Results"),
+        gr.Textbox(lines=10, label="Bug Hunting Results"),
     ],
-    outputs="text",
-    title="Ultimate Code Generator Python: AI Code Generator with Hugging Face Models",
-    description="Select a model and dataset, input a prompt, and generate Python code using AI models."
+    title="GitHub Repository Analyzer",
+    description="Analyze GitHub repositories for optimization suggestions and potential bugs using CodeT5 and other models.",
 )
 
-# Launch the Gradio App
-iface.launch()
+iface.launch()