File size: 2,946 Bytes
1c04ff4
21079e6
0728f6f
21079e6
 
0728f6f
21079e6
 
 
8967b88
 
21079e6
1c04ff4
 
 
0728f6f
21079e6
 
 
 
 
 
0728f6f
21079e6
1c04ff4
21079e6
 
 
1c04ff4
 
 
 
 
 
 
21079e6
 
1c04ff4
21079e6
1c04ff4
21079e6
 
 
0728f6f
1c04ff4
 
 
 
0728f6f
21079e6
 
0728f6f
 
21079e6
0728f6f
21079e6
 
 
 
 
 
 
 
 
0728f6f
21079e6
 
 
 
 
1c04ff4
 
 
21079e6
 
 
 
 
 
 
1c04ff4
21079e6
1c04ff4
21079e6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1c04ff4
 
 
 
 
 
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
import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ---------------------------------
# CONFIG
# ---------------------------------
# BASE_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"
BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
LORA_PATH = "vastu_lora_adapter_975"
DEVICE = "cpu"

HF_TOKEN = os.getenv("HF_TOKEN")  # 🔥 REQUIRED

SYSTEM_PROMPT = """You are a strict and authoritative Vastu Shastra expert.
You clearly classify every placement as IDEAL, ACCEPTABLE, or INADVISABLE.
You always give practical remedies if something is wrong.
Your tone is confident, traditional, and decisive.
Avoid unnecessary philosophy. Be precise and actionable.
"""

# ---------------------------------
# LOAD MODEL (CPU SAFE)
# ---------------------------------
@torch.inference_mode()
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL,
        token=HF_TOKEN,
        trust_remote_code=True,
    )
    tokenizer.pad_token = tokenizer.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        token=HF_TOKEN,                # 🔥 REQUIRED
        torch_dtype=torch.float32,
        device_map="cpu",              # 🔥 FIXED
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )

    model = PeftModel.from_pretrained(
        base_model,
        LORA_PATH,
    )

    model.eval()
    return model, tokenizer


model, tokenizer = load_model()

# ---------------------------------
# GENERATION (CPU OPTIMIZED)
# ---------------------------------
def generate_response(user_prompt):
    prompt = f"""### System:
{SYSTEM_PROMPT}
### User:
{user_prompt}
### Response:
"""

    inputs = tokenizer(prompt, return_tensors="pt")

    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        temperature=0.3,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
    )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return decoded.split("### Response:")[-1].strip()


# ---------------------------------
# GRADIO UI
# ---------------------------------
with gr.Blocks(title="🧭 Vastu AI Advisor (CPU)") as demo:
    gr.Markdown("# 🧭 Vastu AI Advisor")
    gr.Markdown(
        "**CPU-based demo.** Responses may take ~30–45 seconds.\n\n"
        "Ask questions about room placement, directions, and remedies."
    )

    chatbot = gr.Chatbot(height=420)
    msg = gr.Textbox(
        placeholder="Ask a Vastu question (e.g. Is a toilet in NE acceptable?)",
        lines=2,
    )

    def chat(user_msg, history):
        response = generate_response(user_msg)
        history.append((user_msg, response))
        return history, ""

    msg.submit(chat, [msg, chatbot], [chatbot, msg])


demo.launch(
    server_name="0.0.0.0",   # 🔥 REQUIRED FOR HF SPACES
    server_port=7860,
)