Update app.py

app.py

@@ -1,53 +1,56 @@
-from datasets import load_dataset
-from transformers import (
-    AutoModelForSeq2SeqLM,
-    AutoTokenizer,
-    Trainer,
-    DataCollatorForSeq2Seq
-)
-from training_config import training_args
-import os
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
+import gradio as gr
 
-# Load dataset
-dataset = load_dataset("health360/Healix-Shot", split=f"train[:100000]")
+app = FastAPI()
 
 # Initialize model and tokenizer
 model_name = "google/flan-t5-large"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-def tokenize_function(examples):
-    return tokenizer(
-        examples['text'],
-        padding="max_length",
-        truncation=True,
-        max_length=512,
-        return_attention_mask=True
-    )
+class Query(BaseModel):
+    inputs: str
 
-# Process dataset
-train_test_split = dataset.train_test_split(test_size=0.1)
-tokenized_train = train_test_split['train'].map(
-    tokenize_function, 
-    batched=True, 
-    remove_columns=dataset.column_names
-)
-tokenized_eval = train_test_split['test'].map(
-    tokenize_function, 
-    batched=True, 
-    remove_columns=dataset.column_names
-)
+@app.post("/")
+async def generate(query: Query):
+    try:
+        # Tokenize input
+        inputs = tokenizer(query.inputs, return_tensors="pt", max_length=512, truncation=True)
+        
+        # Generate response
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=512,
+            num_beams=4,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.2,
+            early_stopping=True
+        )
+        
+        # Decode response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return {"generated_text": response}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Gradio interface
+def generate_text(prompt):
+    query = Query(inputs=prompt)
+    response = generate(query)
+    return response["generated_text"]
 
-# Initialize trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_train,
-    eval_dataset=tokenized_eval,
-    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
+iface = gr.Interface(
+    fn=generate_text,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your text here..."),
+    outputs="text",
+    title="Medical Assistant",
+    description="Ask me anything about medical topics!"
 )
 
-# Train and save
-trainer.train()
-model.push_to_hub("MjolnirThor/flan-t5-custom-handler")
-tokenizer.push_to_hub("MjolnirThor/flan-t5-custom-handler")
+# Mount the Gradio app
+app = gr.mount_gradio_app(app, iface, path="/")