test4

app.py

@@ -1,97 +1,38 @@
 import gradio as gr
 import torch
-import os
-
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
-
-
-# ------------------------------
-# Authenticate Hugging Face token
-# ------------------------------
-login(token=os.getenv("HF_TOKEN"))
 
-MODEL_ID = "meta-llama/Llama-3.2-1B-Instruct"
-token = os.getenv("HF_TOKEN")
-
-
-# ------------------------------
-# Load tokenizer + model
-# ------------------------------
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID,
-    use_auth_token=token
-)
+MODEL_ID = "WiroAI/WiroAI-Finance-Qwen-1.5B"
 
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float32,      # safer on CPU
-    low_cpu_mem_usage=True,
-    device_map="cpu",
-    use_auth_token=token
+    torch_dtype="auto",
+    device_map="cpu"
 )
 
-
-# ------------------------------
-# System Prompt for FINBOT
-# ------------------------------
 SYSTEM_PROMPT = """
 You are FINBOT, a precise financial assistant.
-
-RULES:
-- NEVER repeat the user's question.
-- NEVER generate extra questions.
-- NEVER explain unless asked.
-- If data is inside <CONTEXT> or <DB_RESULTS>, use ONLY that data.
-- If calculation is required, output ONLY the final number unless the user says “explain”.
-- If asked for analysis, give short bullet points.
-- If data is missing, reply: “Data not available”.
-- Be concise and accurate.
+Follow the rules strictly:
+- Never repeat the question
+- Never explain unless asked
+- Use <CONTEXT> or <DB_RESULTS> only if provided
+- Output only the final number for calculations
 """
 
-
-# ------------------------------
-# Build proper Llama-3 prompt
-# ------------------------------
 def build_prompt(user_msg):
-    return f"""
-<s>[INST] <<SYS>>
-{SYSTEM_PROMPT}
-<</SYS>>
-
-{user_msg}
-[/INST]
-"""
+    return f"{SYSTEM_PROMPT}\n\nUser: {user_msg}\nAssistant:"
 
-
-# ------------------------------
-# Generate function
-# ------------------------------
 def generate(user_input):
     prompt = build_prompt(user_input)
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=100,
-            temperature=0.0,       # deterministic (no randomness)
-            do_sample=False,
-        )
-
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return result.strip()
-
-
-# ------------------------------
-# Gradio UI
-# ------------------------------
-demo = gr.Interface(
-    fn=generate,
-    inputs=gr.Textbox(lines=5, label="Ask FINBOT"),
-    outputs=gr.Textbox(label="FINBOT Answer"),
-    title="FINBOT - Financial Assistant (Llama 3.2 1B)"
-)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=80,
+        temperature=0.0,
+        do_sample=False
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-demo.launch()
+gr.Interface(fn=generate, inputs="text", outputs="text").launch()