something / app.py
anthony01's picture
Update app.py
96988ff verified
import gradio as gr
import os
import shutil
from transformers import AutoTokenizer, AutoModelForCausalLM
# Model details and token
model_name = "Salesforce/codegen-350M-mono"
token = "token" # Replace with your actual token
# Load tokenizer and model with authentication token
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=token)
# Function to explain the Java code
def explain_code(code):
inputs = tokenizer(code, return_tensors="pt")
outputs = model.generate(inputs['input_ids'], max_length=500)
explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
return explanation
# Function to extract Java files from a directory
def extract_java_files(src_dir, dest_dir=None):
java_files = []
for root, dirs, files in os.walk(src_dir):
for file in files:
if file.endswith('.java'):
java_file_path = os.path.join(root, file)
java_files.append(java_file_path)
# If a destination directory is provided, copy the file
if dest_dir:
if not os.path.exists(dest_dir):
os.makedirs(dest_dir)
shutil.copy(java_file_path, os.path.join(dest_dir, file))
return java_files
# Function to read and process uploaded file
def process_uploaded_file(file):
code = file.read().decode("utf-8")
return explain_code(code)
# Create a folder for storing extracted Java files (optional)
src_directory = "path_to_your_directory" # replace with your folder path
dest_directory = "extracted_java_files" # Folder where Java files will be copied
java_files = extract_java_files(src_directory, dest_directory)
# Gradio Interface for selecting uploaded Java file or explaining code
def explain_from_uploaded_or_folder(file=None):
if file:
code = file.read().decode("utf-8")
else:
# If no file uploaded, use the first Java file from the extracted folder (or prompt user)
if java_files:
with open(java_files[0], "r") as f:
code = f.read()
else:
return "No Java files found."
# Process and return explanation of the code
return explain_code(code)
# Create the Gradio interface
iface = gr.Interface(
fn=explain_from_uploaded_or_folder,
inputs=[
gr.File(label="Upload Java File", file_types=[".java"]),
gr.Button("Use First Extracted Java File", visible=not bool(java_files))
],
outputs="text",
live=True
)
# Launch the interface
iface.launch()