Update app.py

app.py

@@ -1,13 +1,33 @@
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 import gradio as gr
 
-# Load your HF model (must be on Hub)
-pipe = pipeline("text-generation", model="eduard76/Llama3-8b-good", trust_remote_code=True)
+model_id = "eduard76/Llama3-8b-good-new"
 
-def chat_fn(message, history):
-    full_input = "\n".join([f"User: {turn[0]}\nAI: {turn[1]}" for turn in history] + [f"User: {message}\nAI:"])
-    result = pipe(full_input, max_new_tokens=256, do_sample=True, temperature=0.7)
-    response = result[0]['generated_text'].split("AI:")[-1].strip()
-    return response
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",  # automatically uses GPU if available
+    torch_dtype=torch.float16,
+    load_in_4bit=True,
+    trust_remote_code=True
+)
+model.eval()
 
-gr.ChatInterface(chat_fn, title="🤖 Chat with Eduard76's LLM").launch()
+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:"
+
+    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
+
+gr.ChatInterface(chat, title="💬 Chat with first Eduard LLM").launch()