rezaenayati commited on
Commit
c646b6b
·
verified ·
1 Parent(s): a9fc148

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +32 -37
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 # Important for ZeroGPU
6
 
7
- # Load models (will be moved to GPU when needed)
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", # ZeroGPU handles this
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 # This decorator is CRITICAL for ZeroGPU
25
  def chat_with_rezAi(messages, history):
26
- conversation = "<|start_header_id|>system<|end_header_id|>\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{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 - fix the max_length parameter
37
  inputs = tokenizer(
38
- conversation,
39
- return_tensors="pt",
40
- truncation=True, # Changed from 'truncate=True'
41
  max_length=2048
42
  )
43
 
44
- # Move inputs to the same device as model
45
  inputs = {k: v.to(model.device) for k, v in inputs.items()}
46
 
47
- # Generate response
48
  with torch.no_grad():
49
  outputs = model.generate(
50
  **inputs,
51
- max_new_tokens=128,
52
- temperature=0.7, # Slightly increased for more variety
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
- # Decode response
60
- response = tokenizer.decode(outputs[0], skip_special_tokens=True)
61
- new_response = response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
62
 
63
- # Clean up response - remove any incomplete tags
 
 
 
 
 
 
 
 
64
  if "<|" in new_response:
65
  new_response = new_response.split("<|")[0].strip()
66
 
67
  return new_response
68
 
69
- # Create Gradio interface
70
  demo = gr.ChatInterface(
71
  fn=chat_with_rezAi,
72
- title="💬 Chat with RezAI",
73
- description="Hi! I'm RezAI, Reza's AI twin. Ask me about his technical background, projects, or experience!",
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__":