Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -2,86 +2,81 @@ import torch
|
|
2 |
from transformers import AutoTokenizer, AutoModelForCausalLM
|
3 |
from peft import PeftModel
|
4 |
import gradio as gr
|
5 |
-
import spaces
|
6 |
|
7 |
-
# Load models
|
8 |
base_model = AutoModelForCausalLM.from_pretrained(
|
9 |
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
|
10 |
torch_dtype=torch.float16,
|
11 |
-
device_map="auto",
|
12 |
trust_remote_code=True
|
13 |
)
|
14 |
|
15 |
tokenizer = AutoTokenizer.from_pretrained("unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit")
|
16 |
|
17 |
-
# Add padding token if missing
|
18 |
if tokenizer.pad_token is None:
|
19 |
tokenizer.pad_token = tokenizer.eos_token
|
20 |
|
21 |
# Load LoRA adapter
|
22 |
model = PeftModel.from_pretrained(base_model, "rezaenayati/RezAi-Model")
|
23 |
|
24 |
-
@spaces.GPU
|
25 |
def chat_with_rezAi(messages, history):
|
26 |
-
|
|
|
27 |
|
28 |
# Add conversation history
|
29 |
for user_msg, assistant_msg in history:
|
30 |
-
conversation += f"<|start_header_id|>user<|end_header_id|>\n{user_msg}<|eot_id|>"
|
31 |
-
conversation += f"<|start_header_id|>assistant<|end_header_id|>\n{assistant_msg}<|eot_id|>"
|
32 |
|
33 |
# Add current message
|
34 |
-
conversation += f"<|start_header_id|>user<|end_header_id|>\n{messages}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
|
35 |
|
36 |
-
# Tokenize
|
37 |
inputs = tokenizer(
|
38 |
-
conversation,
|
39 |
-
return_tensors="pt",
|
40 |
-
truncation=True,
|
41 |
max_length=2048
|
42 |
)
|
43 |
|
44 |
-
# Move
|
45 |
inputs = {k: v.to(model.device) for k, v in inputs.items()}
|
46 |
|
47 |
-
# Generate
|
48 |
with torch.no_grad():
|
49 |
outputs = model.generate(
|
50 |
**inputs,
|
51 |
-
max_new_tokens=
|
52 |
-
temperature=0.
|
53 |
do_sample=True,
|
54 |
pad_token_id=tokenizer.eos_token_id,
|
55 |
-
eos_token_id=tokenizer.eos_token_id
|
56 |
-
repetition_penalty=1.1 # Added to reduce repetition
|
57 |
)
|
58 |
|
59 |
-
#
|
60 |
-
|
61 |
-
new_response = response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
|
62 |
|
63 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
64 |
if "<|" in new_response:
|
65 |
new_response = new_response.split("<|")[0].strip()
|
66 |
|
67 |
return new_response
|
68 |
|
69 |
-
#
|
70 |
demo = gr.ChatInterface(
|
71 |
fn=chat_with_rezAi,
|
72 |
-
title="
|
73 |
-
description="
|
74 |
-
examples=[
|
75 |
-
"Tell me about your background",
|
76 |
-
"What programming languages do you know?",
|
77 |
-
"Walk me through RezAI",
|
78 |
-
"What's your experience with machine learning?",
|
79 |
-
"How did you get into computer science?"
|
80 |
-
],
|
81 |
-
retry_btn=None,
|
82 |
-
undo_btn="Delete Previous",
|
83 |
-
clear_btn="Clear Chat",
|
84 |
-
theme=gr.themes.Soft(), # Added a nice theme
|
85 |
)
|
86 |
|
87 |
if __name__ == "__main__":
|
|
|
2 |
from transformers import AutoTokenizer, AutoModelForCausalLM
|
3 |
from peft import PeftModel
|
4 |
import gradio as gr
|
5 |
+
import spaces
|
6 |
|
7 |
+
# Load models
|
8 |
base_model = AutoModelForCausalLM.from_pretrained(
|
9 |
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
|
10 |
torch_dtype=torch.float16,
|
11 |
+
device_map="auto",
|
12 |
trust_remote_code=True
|
13 |
)
|
14 |
|
15 |
tokenizer = AutoTokenizer.from_pretrained("unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit")
|
16 |
|
|
|
17 |
if tokenizer.pad_token is None:
|
18 |
tokenizer.pad_token = tokenizer.eos_token
|
19 |
|
20 |
# Load LoRA adapter
|
21 |
model = PeftModel.from_pretrained(base_model, "rezaenayati/RezAi-Model")
|
22 |
|
23 |
+
@spaces.GPU
|
24 |
def chat_with_rezAi(messages, history):
|
25 |
+
# Build conversation with proper formatting
|
26 |
+
conversation = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are Reza Enayati, a Computer Science student and entrepreneur from Los Angeles, who is eager to work as a software engineer or machine learning engineer. Answer these questions as if you are in an interview.<|eot_id|>"
|
27 |
|
28 |
# Add conversation history
|
29 |
for user_msg, assistant_msg in history:
|
30 |
+
conversation += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>"
|
31 |
+
conversation += f"<|start_header_id|>assistant<|end_header_id|>\n\n{assistant_msg}<|eot_id|>"
|
32 |
|
33 |
# Add current message
|
34 |
+
conversation += f"<|start_header_id|>user<|end_header_id|>\n\n{messages}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
|
35 |
|
36 |
+
# Tokenize
|
37 |
inputs = tokenizer(
|
38 |
+
conversation,
|
39 |
+
return_tensors="pt",
|
40 |
+
truncation=True,
|
41 |
max_length=2048
|
42 |
)
|
43 |
|
44 |
+
# Move to device
|
45 |
inputs = {k: v.to(model.device) for k, v in inputs.items()}
|
46 |
|
47 |
+
# Generate with higher temperature
|
48 |
with torch.no_grad():
|
49 |
outputs = model.generate(
|
50 |
**inputs,
|
51 |
+
max_new_tokens=150,
|
52 |
+
temperature=0.5, # You asked for 5, but that's too high (0.5 is good)
|
53 |
do_sample=True,
|
54 |
pad_token_id=tokenizer.eos_token_id,
|
55 |
+
eos_token_id=tokenizer.eos_token_id
|
|
|
56 |
)
|
57 |
|
58 |
+
# Extract ONLY the new assistant response
|
59 |
+
full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
60 |
|
61 |
+
# Split by the last assistant header and get only the new response
|
62 |
+
if "<|start_header_id|>assistant<|end_header_id|>" in full_response:
|
63 |
+
response_parts = full_response.split("<|start_header_id|>assistant<|end_header_id|>")
|
64 |
+
new_response = response_parts[-1].strip()
|
65 |
+
else:
|
66 |
+
new_response = full_response.strip()
|
67 |
+
|
68 |
+
# Clean up any remaining special tokens or incomplete parts
|
69 |
+
new_response = new_response.replace("<|eot_id|>", "").strip()
|
70 |
if "<|" in new_response:
|
71 |
new_response = new_response.split("<|")[0].strip()
|
72 |
|
73 |
return new_response
|
74 |
|
75 |
+
# Simple Gradio interface
|
76 |
demo = gr.ChatInterface(
|
77 |
fn=chat_with_rezAi,
|
78 |
+
title="Chat with RezAI",
|
79 |
+
description="Ask me about Reza's background and experience!"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
80 |
)
|
81 |
|
82 |
if __name__ == "__main__":
|