File size: 6,213 Bytes
4f06d80
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import os

# Set up model parameters
MODEL_ID = "alaamostafa/Microsoft-Phi-2"
BASE_MODEL_ID = "microsoft/phi-2"

# Force CPU usage and set up offload directory
device = "cpu"
print(f"Using device: {device}")
os.makedirs("offload_dir", exist_ok=True)

# Disable bitsandbytes for CPU usage
os.environ["BITSANDBYTES_NOWELCOME"] = "1"

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)

# Load base model with simple CPU configuration, avoiding device_map and 8-bit loading
print("Loading base model...")
try:
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.float32,  # Use float32 for CPU
        trust_remote_code=True,
        low_cpu_mem_usage=True,     # Optimize for lower memory usage
        offload_folder="offload_dir" # Set offload directory
    )
    
    # Load the fine-tuned adapter
    print(f"Loading adapter from {MODEL_ID}...")
    model = PeftModel.from_pretrained(
        base_model,
        MODEL_ID,
        offload_folder="offload_dir"
    )
    
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    # Create a placeholder error message for the UI
    error_message = f"Failed to load model: {str(e)}\n\nThis Space may need a GPU to run properly."

def generate_text(
    prompt, 
    max_length=256,  # Reduced for CPU
    temperature=0.7, 
    top_p=0.9, 
    top_k=40, 
    repetition_penalty=1.1
):
    """Generate text based on prompt with the fine-tuned model"""
    try:
        # Prepare input
        inputs = tokenizer(prompt, return_tensors="pt")
        
        # Generate text
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=max_length,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                repetition_penalty=repetition_penalty,
                do_sample=temperature > 0,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode and return the generated text
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return generated_text
    except Exception as e:
        return f"Error generating text: {str(e)}"

# Create the Gradio interface
css = """
.gradio-container {max-width: 800px !important}
.gr-prose code {white-space: pre-wrap !important}
"""

title = "Neuroscience Fine-tuned Phi-2 Model (CPU Version)"
description = """
This is a fine-tuned version of Microsoft's Phi-2 model, adapted specifically for neuroscience domain content.
⚠️ **Note: This model is running on CPU which means responses will be slower.** ⚠️

For best performance:
- Keep your prompts focused and clear
- Use shorter maximum length settings (128-256)
- Be patient as generation can take 30+ seconds

**Example prompts:**
- Recent advances in neuroimaging suggest that
- The role of dopamine in learning and memory involves
- Explain the concept of neuroplasticity in simple terms
- What are the key differences between neurons and glial cells?
"""

# Check if model loaded successfully
if 'error_message' in locals():
    # Simple error interface
    demo = gr.Interface(
        fn=lambda x: error_message,
        inputs=gr.Textbox(label="This model cannot be loaded on CPU"),
        outputs=gr.Textbox(),
        title=title,
        description=description
    )
else:
    # Full interface
    with gr.Blocks(css=css) as demo:
        gr.Markdown(f"# {title}")
        gr.Markdown(description)
        
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Enter your prompt",
                    placeholder="Recent advances in neuroscience suggest that",
                    lines=5
                )
                
                with gr.Row():
                    submit_btn = gr.Button("Generate", variant="primary")
                    clear_btn = gr.Button("Clear")
                
                with gr.Accordion("Advanced Options", open=False):
                    max_length = gr.Slider(
                        minimum=64, maximum=512, value=256, step=64,
                        label="Maximum Length (lower is faster on CPU)"
                    )
                    temperature = gr.Slider(
                        minimum=0.0, maximum=1.5, value=0.7, step=0.1,
                        label="Temperature (0 = deterministic, 0.7 = creative, 1.5 = random)"
                    )
                    top_p = gr.Slider(
                        minimum=0.1, maximum=1.0, value=0.9, step=0.1,
                        label="Top-p (nucleus sampling)"
                    )
                    top_k = gr.Slider(
                        minimum=1, maximum=100, value=40, step=1,
                        label="Top-k"
                    )
                    repetition_penalty = gr.Slider(
                        minimum=1.0, maximum=2.0, value=1.1, step=0.1,
                        label="Repetition Penalty"
                    )
                
            with gr.Column():
                output = gr.Textbox(
                    label="Generated Text",
                    lines=20
                )
        
        # Set up event handlers
        submit_btn.click(
            fn=generate_text,
            inputs=[prompt, max_length, temperature, top_p, top_k, repetition_penalty],
            outputs=output
        )
        clear_btn.click(
            fn=lambda: ("", None),
            inputs=None,
            outputs=[prompt, output]
        )
        
        # Example prompts
        examples = [
            ["Recent advances in neuroimaging suggest that"],
            ["The role of dopamine in learning and memory involves"],
            ["Explain the concept of neuroplasticity in simple terms"],
            ["What are the key differences between neurons and glial cells?"]
        ]
        
        gr.Examples(
            examples=examples,
            inputs=prompt
        )

# Launch the app
demo.launch()