Update app.py

app.py

@@ -4,30 +4,30 @@ import gradio as gr
 
 model_id = "eduard76/Llama3-8b-good-new"
 
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
+
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="auto",  # automatically uses GPU if available
+    device_map="auto",
     torch_dtype=torch.float16,
-    load_in_4bit=True,
+    quantization_config=quant_config,
     trust_remote_code=True
 )
-model.eval()
 
-def chat(user_input, history):
-    history_text = "\n".join([f"User: {u}\nAI: {a}" for u, a in history])
-    prompt = f"{history_text}\nUser: {user_input}\nAI:"
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+def chat(user_input):
+    prompt = f"User: {user_input}\nAI:"
+    response = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]["generated_text"]
+    return response[len(prompt):].strip()
 
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            do_sample=True,
-            temperature=0.001
-        )
-    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = generated.split("AI:")[-1].strip()
-    return answer
+iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Llama3 8B Chat")
 
-gr.ChatInterface(chat, title="💬 Chat with first Eduard LLM").launch()
+if __name__ == "__main__":
+    iface.launch()