Update app.py

app.py

@@ -1,49 +1,54 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import gradio as gr
 
 model_id = "eduard76/Llama3-8b-good-new"
 
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="auto",  # poate fi "cuda:0" sau "cpu" dacă ai eroare
-    torch_dtype=torch.float16,  # sau .bfloat16 dacă vrei
+    device_map="auto",
+    torch_dtype=torch.float16,
     trust_remote_code=True
 )
+model.eval()
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
+# Lista de topicuri acoperite
 covered_topics = {
     "ospf", "bgp", "eigrp", "vxlan", "evpn", "network design", "acl", "routing",
     "spine", "leaf", "underlay", "overlay", "mpls", "qos", "firewall",
     "vpn", "vlan", "subnet", "cidr"
 }
 
+# Funcția principală de chat
 def chat(user_input):
-    prompt = f"""You are a Cisco-certified network architect trained in OSPF, BGP, EIGRP, VLAN, STP, RSTP design principles.
-If the user's question is unclear, clarify first.
-If the topic is outside OSPF, BGP, EIGRP, VLAN, STP, RSTP, respond with: "I'm not trained on that topic."
-Give short, clear, non-repetitive answers.
-
-User: {user_input}
-AI:"""
-
-    response = pipe(
-        prompt,
-        max_new_tokens=256,
-        do_sample=False,
-        temperature=0.0,
-        repetition_penalty=1.2,
-        no_repeat_ngram_size=5,
-        top_k=50,
-        top_p=0.9
-    )[0]["generated_text"]
-
-    return response[len(prompt):].strip()
-
-iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Eduard's 1st virtual Architect")
+    prompt = f"""### Human: {user_input}\n### Assistant:"""
+
+    input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output = model.generate(
+            **input_ids,
+            max_new_tokens=256,
+            do_sample=False,
+            temperature=0.0,
+            repetition_penalty=1.2,
+            no_repeat_ngram_size=5,
+            top_k=50,
+            top_p=0.9
+        )
+
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Scoate promptul inițial din răspuns
+    if "### Assistant:" in response:
+        response = response.split("### Assistant:")[-1].strip()
+
+    return response
+
+# Interfață Gradio
+iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Eduard's Virtual Architect – LLaMA3 Fine-Tuned")
 
 if __name__ == "__main__":
     iface.launch()