Update app.py

app.py

@@ -1,51 +1,59 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+# app.py
+from awq import AutoAWQForCausalLM
+from transformers import AutoTokenizer
 import torch
+import gradio as gr
 
-# Model aur tokenizer load karo
-model_name = "TheBloke/Mistral-7B-v0.1-AWQ"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+# Model name from Hugging Face
+MODEL_NAME = "TheBloke/Mistral-7B-v0.1-AWQ"
 
-def generate_response(prompt, max_tokens=512, temperature=0.7):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        temperature=temperature,
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Gradio respond function
-def respond(message, history, system_message, max_tokens, temperature, top_p, hf_token=None):
-    prompt = system_message + "\n"
-    for h in history:
-        prompt += f"User: {h['user']}\nBot: {h['bot']}\n"
-    prompt += f"User: {message}\nBot: "
-    
-    response = generate_response(prompt, max_tokens=max_tokens, temperature=temperature)
-    yield response
-
-# Gradio ChatInterface
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
-    ],
+# Load the model
+print("🚀 Loading Mistral 7B v0.1 AWQ model...")
+model = AutoAWQForCausalLM.from_quantized(
+    MODEL_NAME,
+    fuse_layers=True,
+    trust_remote_code=False,
+    safetensors=True
 )
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=False)
+print("✅ Model loaded successfully!")
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.Text("Login not required for local model")
-    chatbot.render()
+# Text generation function
+def generate_text(prompt, temperature, max_tokens):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-if __name__ == "__main__":
-    demo.launch()
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+    # Clean the output (remove the original prompt from response)
+    if prompt in response:
+        response = response[len(prompt):].strip()
 
-   
+    return response
+
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(lines=3, placeholder="Ask Mistral something...", label="Prompt"),
+        gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(50, 1024, value=512, step=10, label="Max Tokens")
+    ],
+    outputs=gr.Textbox(lines=10, label="Response"),
+    title="🧠 Mistral 7B v0.1 AWQ",
+    description="Run the quantized Mistral 7B v0.1 model locally or on Google Colab using Gradio.",
+    theme="default"
+)
+
+if __name__ == "__main__":
+    interface.launch(share=True)