Update app.py

app.py

@@ -1,50 +1,52 @@
-   import os
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+ import os
 import gradio as gr
-from huggingface_hub import login
-
-# Log in with HF token
-login(os.environ.get("HF_TOKEN"))
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+import torch
 
-# Model details
-model_name = "mistralai/Mistral-7B-Instruct-v0.3"
-
-# Quantization config
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True
-)
+# Load model from environment variable or fallback
+model_id = os.getenv("MODEL_ID", "TheBloke/Mistral-7B-Instruct-v0.3-GPTQ")
 
 # Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    quantization_config=quantization_config,
+    model_id,
     device_map="auto",
-    low_cpu_mem_usage=True,
-    trust_remote_code=True
+    torch_dtype=torch.float16,
 )
 
-# Chat function
-def chat(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
-    outputs = model.generate(
-        inputs.input_ids,
-        max_length=100,
-        pad_token_id=tokenizer.eos_token_id
+# Initialize streamer for live output
+streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+# Define the chat function
+def chat(message, history):
+    history = history or []
+    conversation = history + [(message, "")]
+
+    prompt = tokenizer.apply_chat_template(
+        conversation,
+        tokenize=False,
+        add_generation_prompt=True
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Gradio UI
-iface = gr.Interface(
-    fn=chat,
-    inputs=gr.Textbox(placeholder="Type your message here..."),
-    outputs="text",
-    title="Vynix AI",
-    description="Chat with Vynix AI — powered by Mistral 7B"
-)
 
-iface.launch()
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.95,
+        )
+
+    decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    reply = decoded.split(message)[-1].strip()
+
+    return reply
+
+# Build Gradio UI
+interface = gr.ChatInterface(fn=chat, title="🧠 Vynix AI")
+
+# Launch App
+if _name_ == "_main_":
+    interface.launch()