Update app.py
Browse files
app.py
CHANGED
@@ -1,148 +1,88 @@
|
|
1 |
-
|
2 |
-
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
token = os.environ["HUB_TOKEN"]
|
18 |
-
device = "cuda" if torch.cuda.is_available() else "cpu"
|
19 |
-
|
20 |
-
PAD_TOKEN = "<|pad|>"
|
21 |
-
EOS_TOKEN = "<|endoftext|>"
|
22 |
-
UNK_TOKEN = "<|unk|>"
|
23 |
-
MAX_INPUT_TOKENS = 1024 # max tokens from context
|
24 |
-
|
25 |
-
|
26 |
-
tokenizer = AutoTokenizer.from_pretrained(REPO, use_auth_token=token, trust_remote_code=True)
|
27 |
-
tokenizer.truncation_side = "left" # ensures if truncate, then keep the last N tokens of the prompt going L -> R
|
28 |
-
|
29 |
-
if device == "cuda":
|
30 |
-
model = AutoModelForCausalLM.from_pretrained(REPO, use_auth_token=token, trust_remote_code=True, low_cpu_mem_usage=True).to(device, dtype=torch.bfloat16)
|
31 |
-
else:
|
32 |
-
model = AutoModelForCausalLM.from_pretrained(REPO, use_auth_token=token, trust_remote_code=True, low_cpu_mem_usage=True)
|
33 |
-
|
34 |
-
model.eval()
|
35 |
-
|
36 |
-
|
37 |
-
custom_css = """
|
38 |
-
.gradio-container {
|
39 |
-
background-color: #0D1525;
|
40 |
-
color:white
|
41 |
-
}
|
42 |
-
#orange-button {
|
43 |
-
background: #F26207 !important;
|
44 |
-
color: white;
|
45 |
-
}
|
46 |
-
.cm-gutters{
|
47 |
-
border: none !important;
|
48 |
-
}
|
49 |
-
"""
|
50 |
-
|
51 |
-
def post_processing(prompt, completion):
|
52 |
-
return prompt + completion
|
53 |
-
# completion = "<span style='color: #499cd5;'>" + completion + "</span>"
|
54 |
-
# prompt = "<span style='color: black;'>" + prompt + "</span>"
|
55 |
-
# code_html = f"<hr><br><pre style='font-size: 14px'><code>{prompt}{completion}</code></pre><br><hr>"
|
56 |
-
# return code_html
|
57 |
-
|
58 |
-
|
59 |
-
def code_generation(prompt, max_new_tokens, temperature=0.2, seed=42, top_p=0.9, top_k=None, use_cache=True, repetition_penalty=1.0):
|
60 |
-
|
61 |
-
# truncates the prompt to MAX_INPUT_TOKENS if its too long
|
62 |
-
x = tokenizer.encode(prompt, return_tensors="pt", max_length=MAX_INPUT_TOKENS, truncation=True).to(device)
|
63 |
-
print("Prompt shape: ", x.shape) # just adding to see in the space logs in prod
|
64 |
-
set_seed(seed)
|
65 |
-
y = model.generate(x,
|
66 |
-
max_new_tokens=max_new_tokens,
|
67 |
-
temperature=temperature,
|
68 |
-
pad_token_id=tokenizer.pad_token_id,
|
69 |
-
eos_token_id=tokenizer.eos_token_id,
|
70 |
-
top_p=top_p,
|
71 |
-
top_k=top_k,
|
72 |
-
use_cache=use_cache,
|
73 |
-
repetition_penalty=repetition_penalty
|
74 |
-
)
|
75 |
-
completion = tokenizer.decode(y[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
76 |
-
completion = completion[len(prompt):]
|
77 |
-
return post_processing(prompt, completion)
|
78 |
-
|
79 |
-
|
80 |
-
demo = gr.Blocks(
|
81 |
-
css=custom_css
|
82 |
)
|
83 |
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
)
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
|
120 |
-
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
from transformers import pipeline
|
3 |
+
|
4 |
+
# App Title
|
5 |
+
st.set_page_config(page_title="ML Assistant with Replit LLM", layout="wide")
|
6 |
+
st.title("🤖 ML Assistant with Replit LLM")
|
7 |
+
st.write("Interact with the Replit LLM for machine learning workflows and AI-driven coding assistance.")
|
8 |
+
|
9 |
+
# Sidebar Configuration
|
10 |
+
st.sidebar.title("Configuration")
|
11 |
+
api_key = st.sidebar.text_input("Replit LLM API Key", type="password")
|
12 |
+
model_name = st.sidebar.text_input("Hugging Face Model Name", "Canstralian/RabbitRedux")
|
13 |
+
task_type = st.sidebar.selectbox(
|
14 |
+
"Choose a Task",
|
15 |
+
["Text Generation", "Pseudocode to Python", "ML Debugging", "Code Optimization"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
16 |
)
|
17 |
|
18 |
+
# Ensure API Key is Provided
|
19 |
+
if not api_key:
|
20 |
+
st.warning("Please provide your Replit LLM API Key in the sidebar to continue.")
|
21 |
+
st.stop()
|
22 |
+
|
23 |
+
# Initialize Replit LLM Pipeline
|
24 |
+
try:
|
25 |
+
nlp_pipeline = pipeline("text2text-generation", model=model_name)
|
26 |
+
st.success("Model loaded successfully!")
|
27 |
+
except Exception as e:
|
28 |
+
st.error(f"Error loading model: {e}")
|
29 |
+
st.stop()
|
30 |
+
|
31 |
+
# Input Section
|
32 |
+
st.subheader("Input Your Query")
|
33 |
+
user_input = st.text_area("Enter your query or task description", height=150)
|
34 |
+
|
35 |
+
# Submit Button
|
36 |
+
if st.button("Generate Output"):
|
37 |
+
if user_input.strip() == "":
|
38 |
+
st.warning("Please enter a valid input.")
|
39 |
+
else:
|
40 |
+
with st.spinner("Processing..."):
|
41 |
+
try:
|
42 |
+
# Generate response using Replit LLM
|
43 |
+
output = nlp_pipeline(user_input)
|
44 |
+
response = output[0]["generated_text"]
|
45 |
+
st.subheader("AI Response")
|
46 |
+
st.write(response)
|
47 |
+
except Exception as e:
|
48 |
+
st.error(f"An error occurred: {e}")
|
49 |
+
|
50 |
+
# Additional ML Features
|
51 |
+
st.subheader("Advanced Machine Learning Assistance")
|
52 |
+
|
53 |
+
if task_type == "Text Generation":
|
54 |
+
st.info("Use the input box to generate text-based output.")
|
55 |
+
elif task_type == "Pseudocode to Python":
|
56 |
+
st.info("Provide pseudocode, and the Replit LLM will attempt to generate Python code.")
|
57 |
+
example = st.button("Show Example")
|
58 |
+
if example:
|
59 |
+
st.code("""
|
60 |
+
# Pseudocode
|
61 |
+
FOR each item IN list:
|
62 |
+
IF item > threshold:
|
63 |
+
PRINT "Above Threshold"
|
64 |
+
|
65 |
+
# Expected Python Output
|
66 |
+
for item in my_list:
|
67 |
+
if item > threshold:
|
68 |
+
print("Above Threshold")
|
69 |
+
""")
|
70 |
+
elif task_type == "ML Debugging":
|
71 |
+
st.info("Describe your ML pipeline error for debugging suggestions.")
|
72 |
+
elif task_type == "Code Optimization":
|
73 |
+
st.info("Paste your Python code for optimization recommendations.")
|
74 |
+
user_code = st.text_area("Paste your Python code", height=200)
|
75 |
+
if st.button("Optimize Code"):
|
76 |
+
with st.spinner("Analyzing and optimizing..."):
|
77 |
+
try:
|
78 |
+
optimization_prompt = f"Optimize the following Python code:\n\n{user_code}"
|
79 |
+
output = nlp_pipeline(optimization_prompt)
|
80 |
+
optimized_code = output[0]["generated_text"]
|
81 |
+
st.subheader("Optimized Code")
|
82 |
+
st.code(optimized_code)
|
83 |
+
except Exception as e:
|
84 |
+
st.error(f"An error occurred: {e}")
|
85 |
+
|
86 |
+
# Footer
|
87 |
+
st.write("---")
|
88 |
+
st.write("Powered by [Replit LLM](https://replit.com) and [Hugging Face](https://huggingface.co).")
|