Update app.py

app.py

@@ -1,36 +1,41 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 
-# Load IndicBART model
-model_name = "ai4bharat/IndicBART"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# Load DeepSeek-R1 model
+tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
 
-def restore_text(input_text, task_type):
-    prefix_map = {
-        "Restore & Correct Tamil Text": "restore: ",
-        "Summarize in Tamil": "summarize: ",
-        "Translate to English": "translate Tamil to English: "
-    }
-    prefix = prefix_map.get(task_type, "restore: ")
+# Device config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-    input_text = prefix + input_text
-    inputs = tokenizer([input_text], return_tensors="pt", padding=True)
-    outputs = model.generate(**inputs, max_length=256, num_beams=4, early_stopping=True)
-    decoded_output = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    return decoded_output[0]
+def restore_tamil_text(prompt, uploaded_file=None):
+    context = ""
+    
+    # If a file is uploaded, read its contents and use it for context
+    if uploaded_file:
+        file_bytes = uploaded_file.read()
+        context = file_bytes.decode("utf-8")
+        full_prompt = f"You are an expert in ancient Tamil. Use the document below to restore and expand the text.\n\nDocument:\n{context}\n\nTask:\n{prompt}"
+    else:
+        full_prompt = f"You are a Tamil literature historian. {prompt}"
+
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7)
+    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return output_text.strip()
 
-# Gradio Interface
 iface = gr.Interface(
-    fn=restore_text,
+    fn=restore_tamil_text,
     inputs=[
-        gr.Textbox(label="Input Tamil Text", lines=8, placeholder="Enter ancient Tamil text here..."),
-        gr.Radio(choices=["Restore & Correct Tamil Text", "Summarize in Tamil", "Translate to English"],
-                 label="Select Task")
+        gr.Textbox(label="Your Prompt (e.g., Restore this ancient Tamil text...)", lines=6),
+        gr.File(label="Optional: Upload Ancient Tamil File (.txt)")
     ],
-    outputs=gr.Textbox(label="Output"),
-    title="🕉️ Ancient Tamil Text Restorer",
-    description="Restores, corrects, and translates 11-12th century Tamil texts using AI4Bharat's IndicBART model."
+    outputs=gr.Textbox(label="Restored Output"),
+    title="🧠 Ancient Tamil Literature AI Agent",
+    description="Hybrid AI using DeepSeek-R1 + optional file context. Restores and expands ancient Tamil literature using DeepSeek LLM."
 )
 
 if __name__ == "__main__":