Update app.py

app.py

@@ -1,29 +1,24 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from memory import update_memory, check_memory
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-with open("persona.txt", "r", encoding="utf-8") as f:
-    personality = f.read()
+# ✅ Model name update karo
+model_name = "TheBloke/Pygmalion-13B-SuperHOT-8K-GPTQ"
 
-model_name = "TheBloke/Pygmalion-7B-GPTQ"
+# ✅ Tokenizer Load karo
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", trust_remote_code=True)
 
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-def chatbot(input_text):
-    memory_response = check_memory(input_text)
-    if memory_response:
-        return memory_response
-
-    prompt = f"{personality}\nUser: {input_text}\nAI:"
-    outputs = generator(prompt, max_length=200, do_sample=True, temperature=0.7, top_p=0.9)
-    response = outputs[0]["generated_text"].split("AI:")[-1].strip()
-
-    update_memory(input_text, response)
-    return response
-
-iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="MoinRomanticBot")
-
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+# ✅ Model Load karo (Automatic Device Selection)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+
+def chat(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    output = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7, top_p=0.9)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+# ✅ Test Chat
+while True:
+    user_input = input("You: ")
+    if user_input.lower() in ["exit", "quit"]:
+        break
+    response = chat(user_input)
+    print("AI:", response)