MistriDevLab / app.py
acecalisto3's picture
Update app.py
c27429b verified
raw
history blame
5.08 kB
import os
import subprocess
import random
import time
from typing import Dict, List, Tuple
from datetime import datetime
import logging
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from huggingface_hub import InferenceClient, cached_download
# --- Configuration ---
VERBOSE = True
MAX_HISTORY = 5
MAX_TOKENS = 2048
TEMPERATURE = 0.7
TOP_P = 0.8
REPETITION_PENALTY = 1.5
MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"
API_KEY = "YOUR_API_KEY"
# --- Logging Setup ---
logging.basicConfig(
filename="app.log",
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
)
# --- Agents ---
agents = [
"WEB_DEV",
"AI_SYSTEM_PROMPT",
"PYTHON_CODE_DEV",
"DATA_SCIENCE",
"UI_UX_DESIGN",
]
# --- Prompts ---
PREFIX = """
{date_time_str}
Purpose: {purpose}
Agent: {agent_name}
"""
LOG_PROMPT = """
PROMPT: {content}
"""
LOG_RESPONSE = """
RESPONSE: {resp}
"""
# --- Functions ---
def format_prompt(message: str, history: List[Tuple[str, str]], max_history_turns: int = 2) -> str:
prompt = ""
for user_prompt, bot_response in history[-max_history_turns:]:
prompt += f"Human: {user_prompt}\nAssistant: {bot_response}\n"
prompt += f"Human: {message}\nAssistant:"
return prompt
def generate(
prompt: str,
history: List[Tuple[str, str]],
agent_name: str = agents[0],
sys_prompt: str = "",
temperature: float = TEMPERATURE,
max_new_tokens: int = MAX_TOKENS,
top_p: float = TOP_P,
repetition_penalty: float = REPETITION_PENALTY,
) -> str:
# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# Create a text generation pipeline
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
# Prepare the full prompt
date_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
full_prompt = PREFIX.format(
date_time_str=date_time_str,
purpose=sys_prompt,
agent_name=agent_name
) + format_prompt(prompt, history)
if VERBOSE:
logging.info(LOG_PROMPT.format(content=full_prompt))
# Generate response
response = generator(
full_prompt,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
repetition_penalty=repetition_penalty,
do_sample=True
)[0]['generated_text']
# Extract the assistant's response
assistant_response = response.split("Assistant:")[-1].strip()
if VERBOSE:
logging.info(LOG_RESPONSE.format(resp=assistant_response))
return assistant_response
def main():
with gr.Blocks() as demo:
gr.Markdown("## FragMixt: The No-Code Development Powerhouse")
gr.Markdown("### Your AI-Powered Development Companion")
# Chat Interface
chatbot = gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel")
# Input Components
message = gr.Textbox(label="Enter your message", placeholder="Ask me anything!")
purpose = gr.Textbox(label="Purpose", placeholder="What is the purpose of this interaction?")
agent_name = gr.Dropdown(label="Agents", choices=[s for s in agents], value=agents[0], interactive=True)
sys_prompt = gr.Textbox(label="System Prompt", max_lines=1, interactive=True)
temperature = gr.Slider(label="Temperature", value=TEMPERATURE, minimum=0.0, maximum=1.0, step=0.05, interactive=True, info="Higher values produce more diverse outputs")
max_new_tokens = gr.Slider(label="Max new tokens", value=MAX_TOKENS, minimum=0, maximum=1048*10, step=64, interactive=True, info="The maximum numbers of new tokens")
top_p = gr.Slider(label="Top-p (nucleus sampling)", value=TOP_P, minimum=0.0, maximum=1, step=0.05, interactive=True, info="Higher values sample more low-probability tokens")
repetition_penalty = gr.Slider(label="Repetition penalty", value=REPETITION_PENALTY, minimum=1.0, maximum=2.0, step=0.05, interactive=True, info="Penalize repeated tokens")
# Button to submit the message
submit_button = gr.Button(value="Send")
# Project Explorer Tab
with gr.Tab("Project Explorer"):
project_path = gr.Textbox(label="Project Path", placeholder="/home/user/app/current_project")
explore_button = gr.Button(value="Explore")
project_output = gr.Textbox(label="File Tree", lines=20)
# Chat App Logic Tab
with gr.Tab("Chat App"):
history = gr.State([])
examples = [
["What is the purpose of this AI agent?", "I am designed to assist with no-code development tasks."],
["Can you help me generate a Python function to calculate the factorial of a number?", "Sure! Here is a Python function to calculate the factorial of a number:"],
["Generate a simple HTML page with a heading and a paragraph.", "